Skip to main content

vtcode_core/models_manager/
model_family.rs

1//! Model family definitions and capability groupings.
2//!
3//! A model family groups models that share certain characteristics like
4//! context windows, supported features, and prompting strategies.
5
6use serde::{Deserialize, Serialize};
7
8use crate::config::models::Provider;
9use crate::config::types::ReasoningEffortLevel;
10
11/// Default context window for most models
12pub const DEFAULT_CONTEXT_WINDOW: i64 = 128_000;
13
14/// Large context window (for models like Gemini)
15pub const LARGE_CONTEXT_WINDOW: i64 = 1_048_576;
16
17/// Medium context window
18pub const MEDIUM_CONTEXT_WINDOW: i64 = 200_000;
19
20/// Shell tool type preference for a model family
21#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]
22pub enum ShellToolType {
23    /// Use default shell tool behavior
24    #[default]
25    Default,
26    /// Use shell command tool
27    ShellCommand,
28    /// Use local shell execution
29    Local,
30    /// Use unified exec pattern (Codex-style)
31    UnifiedExec,
32}
33
34/// Truncation policy for model output
35#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
36pub enum TruncationPolicy {
37    /// Truncate by byte count
38    Bytes(usize),
39    /// Truncate by token count
40    Tokens(usize),
41    /// No truncation
42    None,
43}
44
45impl Default for TruncationPolicy {
46    fn default() -> Self {
47        TruncationPolicy::Bytes(10_000)
48    }
49}
50
51/// A model family groups models that share certain characteristics.
52#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
53pub struct ModelFamily {
54    /// The full model slug used to derive this model family
55    pub slug: String,
56
57    /// The model family name (e.g., "gemini-2.5", "claude-opus")
58    pub family: String,
59
60    /// The provider this model belongs to
61    pub provider: Provider,
62
63    /// Maximum supported context window, if known
64    pub context_window: Option<i64>,
65
66    /// Token threshold for automatic compaction
67    pub auto_compact_token_limit: Option<i64>,
68
69    /// Whether the model supports reasoning summaries
70    pub supports_reasoning_summaries: bool,
71
72    /// Default reasoning effort for this model family
73    pub default_reasoning_effort: Option<ReasoningEffortLevel>,
74
75    /// Whether this model supports parallel tool calls
76    pub supports_parallel_tool_calls: bool,
77
78    /// Whether the model needs special apply_patch instructions
79    pub needs_special_apply_patch_instructions: bool,
80
81    /// Preferred shell tool type for this model family
82    pub shell_type: ShellToolType,
83
84    /// Truncation policy for model output
85    pub truncation_policy: TruncationPolicy,
86
87    /// Names of experimental tools supported by this model family
88    pub experimental_supported_tools: Vec<String>,
89
90    /// Percentage of context window considered usable for inputs
91    pub effective_context_window_percent: i64,
92
93    /// Whether the model supports verbosity settings
94    pub support_verbosity: bool,
95
96    /// Whether the model supports tool use
97    pub supports_tool_use: bool,
98
99    /// Whether the model supports streaming
100    pub supports_streaming: bool,
101
102    /// Whether the model supports thinking/reasoning output
103    pub supports_thinking: bool,
104}
105
106impl Default for ModelFamily {
107    fn default() -> Self {
108        Self {
109            slug: String::new(),
110            family: String::new(),
111            provider: Provider::default(),
112            context_window: Some(DEFAULT_CONTEXT_WINDOW),
113            auto_compact_token_limit: None,
114            supports_reasoning_summaries: false,
115            default_reasoning_effort: None,
116            supports_parallel_tool_calls: false,
117            needs_special_apply_patch_instructions: false,
118            shell_type: ShellToolType::Default,
119            truncation_policy: TruncationPolicy::default(),
120            experimental_supported_tools: Vec::new(),
121            effective_context_window_percent: 95,
122            support_verbosity: false,
123            supports_tool_use: true,
124            supports_streaming: true,
125            supports_thinking: false,
126        }
127    }
128}
129
130impl ModelFamily {
131    /// Create a new model family with the given slug
132    pub fn new(slug: impl Into<String>, family: impl Into<String>, provider: Provider) -> Self {
133        Self {
134            slug: slug.into(),
135            family: family.into(),
136            provider,
137            ..Default::default()
138        }
139    }
140
141    /// Get the auto-compact token limit, computing a default if not set
142    pub fn auto_compact_token_limit(&self) -> Option<i64> {
143        self.auto_compact_token_limit
144            .or(self.context_window.map(Self::default_auto_compact_limit))
145    }
146
147    /// Compute the default auto-compact limit (90% of context window)
148    const fn default_auto_compact_limit(context_window: i64) -> i64 {
149        (context_window * 9) / 10
150    }
151
152    /// Get the model slug
153    pub fn get_model_slug(&self) -> &str {
154        &self.slug
155    }
156
157    /// Check if this family supports a specific feature
158    pub fn supports_feature(&self, feature: &str) -> bool {
159        match feature {
160            "reasoning" | "thinking" => self.supports_thinking,
161            "tool_use" | "tools" => self.supports_tool_use,
162            "streaming" => self.supports_streaming,
163            "parallel_tools" => self.supports_parallel_tool_calls,
164            _ => self
165                .experimental_supported_tools
166                .contains(&feature.to_string()),
167        }
168    }
169}
170
171/// Macro to simplify model family definitions
172#[macro_export]
173macro_rules! model_family {
174    (
175        $slug:expr, $family:expr, $provider:expr $(, $key:ident : $value:expr )* $(,)?
176    ) => {{
177        let mut mf = $crate::models_manager::ModelFamily::new($slug, $family, $provider);
178        $(
179            mf.$key = $value;
180        )*
181        mf
182    }};
183}
184
185/// Internal helper that returns a `ModelFamily` for the given model slug.
186pub fn find_family_for_model(slug: &str) -> ModelFamily {
187    if let Some((provider, raw_slug)) = opencode_provider_and_raw_slug(slug) {
188        let mut family = find_family_for_model(raw_slug);
189        family.slug = slug.to_string();
190        family.provider = provider;
191        return family;
192    }
193
194    // Gemini models
195    if slug.starts_with("gemini-3") {
196        return model_family!(
197            slug, "gemini-3", Provider::Gemini,
198            context_window: Some(LARGE_CONTEXT_WINDOW),
199            supports_thinking: true,
200            supports_parallel_tool_calls: true,
201            supports_reasoning_summaries: true,
202        );
203    }
204    if slug.starts_with("gemini") {
205        return model_family!(
206            slug, "gemini", Provider::Gemini,
207            context_window: Some(LARGE_CONTEXT_WINDOW),
208        );
209    }
210
211    // OpenAI models
212    if slug.starts_with("gpt-5") {
213        return model_family!(
214            slug, "gpt-5", Provider::OpenAI,
215            context_window: Some(DEFAULT_CONTEXT_WINDOW),
216            supports_thinking: true,
217            supports_parallel_tool_calls: true,
218        );
219    }
220    if slug.starts_with("codex") {
221        return model_family!(
222            slug, "codex", Provider::OpenAI,
223            context_window: Some(MEDIUM_CONTEXT_WINDOW),
224            supports_thinking: true,
225            shell_type: ShellToolType::UnifiedExec,
226        );
227    }
228    if slug.starts_with("gpt-oss") || slug.contains("gpt-oss") {
229        return model_family!(
230            slug, "gpt-oss", Provider::OpenAI,
231            context_window: Some(96_000),
232        );
233    }
234    if slug.starts_with("o3") || slug.starts_with("o4") {
235        return model_family!(
236            slug, "o-series", Provider::OpenAI,
237            context_window: Some(MEDIUM_CONTEXT_WINDOW),
238            supports_thinking: true,
239            supports_reasoning_summaries: true,
240            needs_special_apply_patch_instructions: true,
241        );
242    }
243
244    // Anthropic models
245    if slug.starts_with("claude-opus") || slug.contains("opus") {
246        return model_family!(
247            slug, "claude-opus", Provider::Anthropic,
248            context_window: Some(MEDIUM_CONTEXT_WINDOW),
249            supports_thinking: true,
250            supports_parallel_tool_calls: true,
251        );
252    }
253    if slug.starts_with("claude-sonnet") || slug.contains("sonnet") {
254        return model_family!(
255            slug, "claude-sonnet", Provider::Anthropic,
256            context_window: Some(MEDIUM_CONTEXT_WINDOW),
257            supports_thinking: true,
258        );
259    }
260    if slug.starts_with("claude-haiku") || slug.contains("haiku") {
261        return model_family!(
262            slug, "claude-haiku", Provider::Anthropic,
263            context_window: Some(MEDIUM_CONTEXT_WINDOW),
264        );
265    }
266    if slug.starts_with("claude") {
267        return model_family!(
268            slug, "claude", Provider::Anthropic,
269            context_window: Some(MEDIUM_CONTEXT_WINDOW),
270        );
271    }
272
273    // DeepSeek models
274    if slug.contains("deepseek") && slug.contains("reason") {
275        return model_family!(
276            slug, "deepseek-reasoner", Provider::DeepSeek,
277            context_window: Some(DEFAULT_CONTEXT_WINDOW),
278            supports_thinking: true,
279        );
280    }
281    if slug.contains("deepseek") {
282        return model_family!(
283            slug, "deepseek", Provider::DeepSeek,
284            context_window: Some(DEFAULT_CONTEXT_WINDOW),
285        );
286    }
287
288    // Z.AI GLM models
289    if slug.contains("glm-5") {
290        return model_family!(
291            slug, "glm-5", Provider::ZAI,
292            context_window: Some(DEFAULT_CONTEXT_WINDOW),
293            supports_thinking: true,
294        );
295    }
296    if slug.contains("glm") {
297        return model_family!(
298            slug, "glm", Provider::ZAI,
299            context_window: Some(DEFAULT_CONTEXT_WINDOW),
300        );
301    }
302
303    // MiniMax models
304    if slug.contains("minimax") {
305        return model_family!(
306            slug, "minimax", Provider::Minimax,
307            context_window: Some(DEFAULT_CONTEXT_WINDOW),
308            supports_thinking: true,
309        );
310    }
311
312    // Moonshot/Kimi models
313    if slug.contains("kimi") || slug.contains("moonshot") {
314        return model_family!(
315            slug, "kimi", Provider::Moonshot,
316            context_window: Some(DEFAULT_CONTEXT_WINDOW),
317            supports_thinking: slug.contains("thinking"),
318        );
319    }
320
321    // Qwen models (via OpenRouter or Ollama)
322    if slug.contains("qwen") {
323        return model_family!(
324            slug, "qwen", Provider::OpenRouter,
325            context_window: Some(DEFAULT_CONTEXT_WINDOW),
326            supports_thinking: slug.contains("thinking"),
327        );
328    }
329
330    // Ollama local models
331    if slug.starts_with("ollama/") || slug.contains(":") {
332        return model_family!(
333            slug, "ollama-local", Provider::Ollama,
334            context_window: Some(DEFAULT_CONTEXT_WINDOW),
335        );
336    }
337
338    // OpenRouter models (fallback for unrecognized patterns)
339    if slug.contains("/") {
340        return model_family!(
341            slug, "openrouter", Provider::OpenRouter,
342            context_window: Some(DEFAULT_CONTEXT_WINDOW),
343        );
344    }
345
346    // Default fallback
347    model_family!(
348        slug, "unknown", Provider::default(),
349        context_window: Some(DEFAULT_CONTEXT_WINDOW),
350    )
351}
352
353fn opencode_provider_and_raw_slug(slug: &str) -> Option<(Provider, &str)> {
354    if let Some(raw_slug) = slug.strip_prefix("opencode-go/") {
355        Some((Provider::OpenCodeGo, raw_slug))
356    } else if let Some(raw_slug) = slug
357        .strip_prefix("opencode/")
358        .or_else(|| slug.strip_prefix("opencode-zen/"))
359    {
360        Some((Provider::OpenCodeZen, raw_slug))
361    } else {
362        None
363    }
364}
365
366#[cfg(test)]
367mod tests {
368    use super::*;
369
370    #[test]
371    fn test_gemini_family_detection() {
372        let family = find_family_for_model("gemini-3-flash-preview");
373        assert_eq!(family.family, "gemini-3");
374        assert_eq!(family.provider, Provider::Gemini);
375        assert!(family.context_window.unwrap() >= LARGE_CONTEXT_WINDOW);
376    }
377
378    #[test]
379    fn test_gpt5_family_detection() {
380        let family = find_family_for_model("gpt-5.3-codex");
381        assert_eq!(family.family, "gpt-5");
382        assert_eq!(family.provider, Provider::OpenAI);
383        assert!(family.supports_thinking);
384    }
385
386    #[test]
387    fn test_claude_family_detection() {
388        let family = find_family_for_model("claude-opus-4.5");
389        assert_eq!(family.family, "claude-opus");
390        assert_eq!(family.provider, Provider::Anthropic);
391    }
392
393    #[test]
394    fn test_opencode_zen_family_detection_preserves_provider() {
395        let family = find_family_for_model("opencode/gpt-5.4");
396        assert_eq!(family.family, "gpt-5");
397        assert_eq!(family.provider, Provider::OpenCodeZen);
398        assert!(family.supports_thinking);
399    }
400
401    #[test]
402    fn test_opencode_go_family_detection_preserves_provider() {
403        let family = find_family_for_model("opencode-go/kimi-k2.5");
404        assert_eq!(family.family, "kimi");
405        assert_eq!(family.provider, Provider::OpenCodeGo);
406    }
407
408    #[test]
409    fn test_auto_compact_limit() {
410        let family = ModelFamily {
411            context_window: Some(100_000),
412            ..Default::default()
413        };
414        assert_eq!(family.auto_compact_token_limit(), Some(90_000));
415    }
416
417    #[test]
418    fn test_supports_feature() {
419        let family = ModelFamily {
420            supports_thinking: true,
421            supports_tool_use: true,
422            ..Default::default()
423        };
424        assert!(family.supports_feature("thinking"));
425        assert!(family.supports_feature("tool_use"));
426        assert!(!family.supports_feature("unknown"));
427    }
428}