Skip to main content

vtcode_commons/
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::provider::Provider;
9use crate::reasoning::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 Codex exec_command pattern (Codex-style)
31    ExecCommand,
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 the 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.experimental_supported_tools.contains(&feature.to_string()),
165        }
166    }
167}
168
169/// Macro to simplify model family definitions
170#[macro_export]
171macro_rules! model_family {
172    (
173        $slug:expr, $family:expr, $provider:expr $(, $key:ident : $value:expr )* $(,)?
174    ) => {{
175        let mut mf = $crate::model_family::ModelFamily::new($slug, $family, $provider);
176        $(
177            mf.$key = $value;
178        )*
179        mf
180    }};
181}
182
183/// Internal helper that returns a `ModelFamily` for the given model slug.
184pub fn find_family_for_model(slug: &str) -> ModelFamily {
185    if let Some((provider, raw_slug)) = opencode_provider_and_raw_slug(slug) {
186        let mut family = find_family_for_model(raw_slug);
187        family.slug = slug.to_string();
188        family.provider = provider;
189        return family;
190    }
191
192    // Gemini models
193    if slug.starts_with("gemini-3") {
194        return model_family!(
195            slug, "gemini-3", Provider::Gemini,
196            context_window: Some(LARGE_CONTEXT_WINDOW),
197            supports_thinking: true,
198            supports_parallel_tool_calls: true,
199            supports_reasoning_summaries: true,
200        );
201    }
202    if slug.starts_with("gemini") {
203        return model_family!(
204            slug, "gemini", Provider::Gemini,
205            context_window: Some(LARGE_CONTEXT_WINDOW),
206        );
207    }
208
209    // OpenAI models
210    if slug.starts_with("gpt-5") {
211        return model_family!(
212            slug, "gpt-5", Provider::OpenAI,
213            context_window: Some(DEFAULT_CONTEXT_WINDOW),
214            supports_thinking: true,
215            supports_parallel_tool_calls: true,
216        );
217    }
218    if slug.starts_with("codex") {
219        return model_family!(
220            slug, "codex", Provider::OpenAI,
221            context_window: Some(MEDIUM_CONTEXT_WINDOW),
222            supports_thinking: true,
223            shell_type: ShellToolType::ExecCommand,
224        );
225    }
226    if slug.starts_with("gpt-oss") || slug.contains("gpt-oss") {
227        return model_family!(
228            slug, "gpt-oss", Provider::OpenAI,
229            context_window: Some(96_000),
230        );
231    }
232    if slug.starts_with("o3") || slug.starts_with("o4") {
233        return model_family!(
234            slug, "o-series", Provider::OpenAI,
235            context_window: Some(MEDIUM_CONTEXT_WINDOW),
236            supports_thinking: true,
237            supports_reasoning_summaries: true,
238            needs_special_apply_patch_instructions: true,
239        );
240    }
241
242    // Anthropic models
243    if slug.starts_with("claude-opus") || slug.contains("opus") {
244        return model_family!(
245            slug, "claude-opus", Provider::Anthropic,
246            context_window: Some(MEDIUM_CONTEXT_WINDOW),
247            supports_thinking: true,
248            supports_parallel_tool_calls: true,
249        );
250    }
251    if slug.starts_with("claude-sonnet") || slug.contains("sonnet") {
252        return model_family!(
253            slug, "claude-sonnet", Provider::Anthropic,
254            context_window: Some(MEDIUM_CONTEXT_WINDOW),
255            supports_thinking: true,
256        );
257    }
258    if slug.starts_with("claude-haiku") || slug.contains("haiku") {
259        return model_family!(
260            slug, "claude-haiku", Provider::Anthropic,
261            context_window: Some(MEDIUM_CONTEXT_WINDOW),
262        );
263    }
264    if slug.starts_with("claude") {
265        return model_family!(
266            slug, "claude", Provider::Anthropic,
267            context_window: Some(MEDIUM_CONTEXT_WINDOW),
268        );
269    }
270
271    // DeepSeek models
272    if slug.contains("deepseek") && slug.contains("reason") {
273        return model_family!(
274            slug, "deepseek-reasoner", Provider::DeepSeek,
275            context_window: Some(DEFAULT_CONTEXT_WINDOW),
276            supports_thinking: true,
277        );
278    }
279    if slug.contains("deepseek") {
280        return model_family!(
281            slug, "deepseek", Provider::DeepSeek,
282            context_window: Some(DEFAULT_CONTEXT_WINDOW),
283        );
284    }
285
286    // Z.AI GLM models
287    if slug.contains("glm-5") {
288        return model_family!(
289            slug, "glm-5", Provider::ZAI,
290            context_window: Some(DEFAULT_CONTEXT_WINDOW),
291            supports_thinking: true,
292        );
293    }
294    if slug.contains("glm") {
295        return model_family!(
296            slug, "glm", Provider::ZAI,
297            context_window: Some(DEFAULT_CONTEXT_WINDOW),
298        );
299    }
300
301    // MiniMax models
302    if slug.contains("minimax") {
303        return model_family!(
304            slug, "minimax", Provider::Minimax,
305            context_window: Some(DEFAULT_CONTEXT_WINDOW),
306            supports_thinking: true,
307        );
308    }
309
310    // Moonshot/Kimi models
311    if slug.contains("kimi") || slug.contains("moonshot") {
312        return model_family!(
313            slug, "kimi", Provider::Moonshot,
314            context_window: Some(DEFAULT_CONTEXT_WINDOW),
315            supports_thinking: slug.contains("thinking"),
316        );
317    }
318
319    // Qwen models (via OpenRouter or Ollama)
320    if slug.contains("qwen") {
321        return model_family!(
322            slug, "qwen", Provider::OpenRouter,
323            context_window: Some(DEFAULT_CONTEXT_WINDOW),
324            supports_thinking: slug.contains("thinking"),
325        );
326    }
327
328    // Ollama local models
329    if slug.starts_with("ollama/") || slug.contains(":") {
330        return model_family!(
331            slug, "ollama-local", Provider::Ollama,
332            context_window: Some(DEFAULT_CONTEXT_WINDOW),
333        );
334    }
335
336    // OpenRouter models (fallback for unrecognized patterns)
337    if slug.contains("/") {
338        return model_family!(
339            slug, "openrouter", Provider::OpenRouter,
340            context_window: Some(DEFAULT_CONTEXT_WINDOW),
341        );
342    }
343
344    // Default fallback
345    model_family!(
346        slug, "unknown", Provider::default(),
347        context_window: Some(DEFAULT_CONTEXT_WINDOW),
348    )
349}
350
351fn opencode_provider_and_raw_slug(slug: &str) -> Option<(Provider, &str)> {
352    if let Some(raw_slug) = slug.strip_prefix("opencode-go/") {
353        Some((Provider::OpenCodeGo, raw_slug))
354    } else if let Some(raw_slug) =
355        slug.strip_prefix("opencode/").or_else(|| slug.strip_prefix("opencode-zen/"))
356    {
357        Some((Provider::OpenCodeZen, raw_slug))
358    } else {
359        None
360    }
361}
362
363#[cfg(test)]
364mod tests {
365    use super::*;
366
367    #[test]
368    fn test_gemini_family_detection() {
369        let family = find_family_for_model("gemini-3-flash-preview");
370        assert_eq!(family.family, "gemini-3");
371        assert_eq!(family.provider, Provider::Gemini);
372        assert!(family.context_window.unwrap() >= LARGE_CONTEXT_WINDOW);
373    }
374
375    #[test]
376    fn test_gpt5_family_detection() {
377        let family = find_family_for_model("gpt-5.3-codex");
378        assert_eq!(family.family, "gpt-5");
379        assert_eq!(family.provider, Provider::OpenAI);
380        assert!(family.supports_thinking);
381    }
382
383    #[test]
384    fn test_claude_family_detection() {
385        let family = find_family_for_model("claude-opus-4.5");
386        assert_eq!(family.family, "claude-opus");
387        assert_eq!(family.provider, Provider::Anthropic);
388    }
389
390    #[test]
391    fn test_opencode_zen_family_detection_preserves_provider() {
392        let family = find_family_for_model("opencode/gpt-5.4");
393        assert_eq!(family.family, "gpt-5");
394        assert_eq!(family.provider, Provider::OpenCodeZen);
395        assert!(family.supports_thinking);
396    }
397
398    #[test]
399    fn test_opencode_go_family_detection_preserves_provider() {
400        let family = find_family_for_model("opencode-go/kimi-k2.5");
401        assert_eq!(family.family, "kimi");
402        assert_eq!(family.provider, Provider::OpenCodeGo);
403    }
404
405    #[test]
406    fn test_auto_compact_limit() {
407        let family = ModelFamily {
408            context_window: Some(100_000),
409            ..Default::default()
410        };
411        assert_eq!(family.auto_compact_token_limit(), Some(90_000));
412    }
413
414    #[test]
415    fn test_supports_feature() {
416        let family = ModelFamily {
417            supports_thinking: true,
418            supports_tool_use: true,
419            ..Default::default()
420        };
421        assert!(family.supports_feature("thinking"));
422        assert!(family.supports_feature("tool_use"));
423        assert!(!family.supports_feature("unknown"));
424    }
425}