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
12const DEFAULT_CONTEXT_WINDOW: i64 = 128_000;
13
14/// Large context window (for models like Gemini)
15const LARGE_CONTEXT_WINDOW: i64 = 1_048_576;
16
17/// Medium context window
18const 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    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    context_window: Option<i64>,
65
66    /// Token threshold for automatic compaction
67    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    default_reasoning_effort: Option<ReasoningEffortLevel>,
74
75    /// Whether the model supports parallel tool calls
76    supports_parallel_tool_calls: bool,
77
78    /// Whether the model needs special apply_patch instructions
79    needs_special_apply_patch_instructions: bool,
80
81    /// Preferred shell tool type for this model family
82    shell_type: ShellToolType,
83
84    /// Truncation policy for model output
85    truncation_policy: TruncationPolicy,
86
87    /// Names of experimental tools supported by this model family
88    experimental_supported_tools: Vec<String>,
89
90    /// Percentage of context window considered usable for inputs
91    effective_context_window_percent: i64,
92
93    /// Whether the model supports verbosity settings
94    support_verbosity: bool,
95
96    /// Whether the model supports tool use
97    supports_tool_use: bool,
98
99    /// Whether the model supports streaming
100    supports_streaming: bool,
101
102    /// Whether the model supports thinking/reasoning output
103    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    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    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    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-6") {
211        return model_family!(
212            slug, "gpt-6", Provider::OpenAI,
213            context_window: Some(LARGE_CONTEXT_WINDOW),
214            supports_thinking: true,
215            supports_parallel_tool_calls: true,
216        );
217    }
218    if slug.starts_with("gpt-5") {
219        return model_family!(
220            slug, "gpt-5", Provider::OpenAI,
221            context_window: Some(DEFAULT_CONTEXT_WINDOW),
222            supports_thinking: true,
223            supports_parallel_tool_calls: true,
224        );
225    }
226    if slug.starts_with("codex") {
227        return model_family!(
228            slug, "codex", Provider::OpenAI,
229            context_window: Some(MEDIUM_CONTEXT_WINDOW),
230            supports_thinking: true,
231            shell_type: ShellToolType::ExecCommand,
232        );
233    }
234    if slug.starts_with("gpt-oss") || slug.contains("gpt-oss") {
235        return model_family!(
236            slug, "gpt-oss", Provider::OpenAI,
237            context_window: Some(96_000),
238        );
239    }
240    if slug.starts_with("o3") || slug.starts_with("o4") {
241        return model_family!(
242            slug, "o-series", Provider::OpenAI,
243            context_window: Some(MEDIUM_CONTEXT_WINDOW),
244            supports_thinking: true,
245            supports_reasoning_summaries: true,
246            needs_special_apply_patch_instructions: true,
247        );
248    }
249
250    // Anthropic models
251    if slug.starts_with("claude-opus") || slug.contains("opus") {
252        return model_family!(
253            slug, "claude-opus", Provider::Anthropic,
254            context_window: Some(MEDIUM_CONTEXT_WINDOW),
255            supports_thinking: true,
256            supports_parallel_tool_calls: true,
257        );
258    }
259    if slug.starts_with("claude-sonnet") || slug.contains("sonnet") {
260        return model_family!(
261            slug, "claude-sonnet", Provider::Anthropic,
262            context_window: Some(MEDIUM_CONTEXT_WINDOW),
263            supports_thinking: true,
264        );
265    }
266    if slug.starts_with("claude-haiku") || slug.contains("haiku") {
267        return model_family!(
268            slug, "claude-haiku", Provider::Anthropic,
269            context_window: Some(MEDIUM_CONTEXT_WINDOW),
270        );
271    }
272    if slug.starts_with("claude") {
273        return model_family!(
274            slug, "claude", Provider::Anthropic,
275            context_window: Some(MEDIUM_CONTEXT_WINDOW),
276        );
277    }
278
279    // DeepSeek models
280    if slug.contains("deepseek") && slug.contains("reason") {
281        return model_family!(
282            slug, "deepseek-reasoner", Provider::DeepSeek,
283            context_window: Some(DEFAULT_CONTEXT_WINDOW),
284            supports_thinking: true,
285        );
286    }
287    if slug.contains("deepseek") {
288        return model_family!(
289            slug, "deepseek", Provider::DeepSeek,
290            context_window: Some(DEFAULT_CONTEXT_WINDOW),
291        );
292    }
293
294    // Official Meta AI Muse models
295    if slug.starts_with("muse-spark-") {
296        return model_family!(
297            slug, "muse-spark", Provider::Meta,
298            context_window: Some(LARGE_CONTEXT_WINDOW),
299            supports_thinking: true,
300            supports_parallel_tool_calls: true,
301            supports_reasoning_summaries: false,
302        );
303    }
304
305    // Z.AI GLM models
306    if slug.contains("glm-5") {
307        return model_family!(
308            slug, "glm-5", Provider::ZAI,
309            context_window: Some(DEFAULT_CONTEXT_WINDOW),
310            supports_thinking: true,
311        );
312    }
313    if slug.contains("glm") {
314        return model_family!(
315            slug, "glm", Provider::ZAI,
316            context_window: Some(DEFAULT_CONTEXT_WINDOW),
317        );
318    }
319
320    // MiniMax models
321    if slug.contains("minimax") {
322        return model_family!(
323            slug, "minimax", Provider::Minimax,
324            context_window: Some(DEFAULT_CONTEXT_WINDOW),
325            supports_thinking: true,
326        );
327    }
328
329    // Moonshot/Kimi models
330    if slug.contains("kimi") || slug.contains("moonshot") {
331        return model_family!(
332            slug, "kimi", Provider::Moonshot,
333            context_window: Some(DEFAULT_CONTEXT_WINDOW),
334            supports_thinking: slug.contains("thinking"),
335        );
336    }
337
338    // Qwen models (via OpenRouter or Ollama)
339    if slug.contains("qwen") {
340        return model_family!(
341            slug, "qwen", Provider::OpenRouter,
342            context_window: Some(DEFAULT_CONTEXT_WINDOW),
343            supports_thinking: slug.contains("thinking"),
344        );
345    }
346
347    // Ollama local models
348    if slug.starts_with("ollama/") || slug.contains(":") {
349        return model_family!(
350            slug, "ollama-local", Provider::Ollama,
351            context_window: Some(DEFAULT_CONTEXT_WINDOW),
352        );
353    }
354
355    // OpenRouter models (fallback for unrecognized patterns)
356    if slug.contains("/") {
357        return model_family!(
358            slug, "openrouter", Provider::OpenRouter,
359            context_window: Some(DEFAULT_CONTEXT_WINDOW),
360        );
361    }
362
363    // Default fallback
364    model_family!(
365        slug, "unknown", Provider::default(),
366        context_window: Some(DEFAULT_CONTEXT_WINDOW),
367    )
368}
369
370fn opencode_provider_and_raw_slug(slug: &str) -> Option<(Provider, &str)> {
371    if let Some(raw_slug) = slug.strip_prefix("opencode-go/") {
372        Some((Provider::OpenCodeGo, raw_slug))
373    } else if let Some(raw_slug) = slug.strip_prefix("opencode/").or_else(|| slug.strip_prefix("opencode-zen/")) {
374        Some((Provider::OpenCodeZen, raw_slug))
375    } else {
376        None
377    }
378}
379
380#[cfg(test)]
381mod tests {
382    use super::*;
383
384    #[test]
385    fn test_gemini_family_detection() {
386        let family = find_family_for_model("gemini-3-flash-preview");
387        assert_eq!(family.family, "gemini-3");
388        assert_eq!(family.provider, Provider::Gemini);
389        assert!(family.context_window.unwrap() >= LARGE_CONTEXT_WINDOW);
390    }
391
392    #[test]
393    fn test_gpt5_family_detection() {
394        let family = find_family_for_model("gpt-5-codex");
395        assert_eq!(family.family, "gpt-5");
396        assert_eq!(family.provider, Provider::OpenAI);
397        assert!(family.supports_thinking);
398    }
399
400    #[test]
401    fn test_claude_family_detection() {
402        let family = find_family_for_model("claude-opus-4.5");
403        assert_eq!(family.family, "claude-opus");
404        assert_eq!(family.provider, Provider::Anthropic);
405    }
406
407    #[test]
408    fn test_opencode_zen_family_detection_preserves_provider() {
409        let family = find_family_for_model("opencode/gpt-5.4");
410        assert_eq!(family.family, "gpt-5");
411        assert_eq!(family.provider, Provider::OpenCodeZen);
412        assert!(family.supports_thinking);
413    }
414
415    #[test]
416    fn test_opencode_go_family_detection_preserves_provider() {
417        let family = find_family_for_model("opencode-go/kimi-k2.5");
418        assert_eq!(family.family, "kimi");
419        assert_eq!(family.provider, Provider::OpenCodeGo);
420    }
421
422    #[test]
423    fn test_auto_compact_limit() {
424        let family = ModelFamily {
425            context_window: Some(100_000),
426            ..Default::default()
427        };
428        assert_eq!(family.auto_compact_token_limit(), Some(90_000));
429    }
430
431    #[test]
432    fn test_supports_feature() {
433        let family = ModelFamily {
434            supports_thinking: true,
435            supports_tool_use: true,
436            ..Default::default()
437        };
438        assert!(family.supports_feature("thinking"));
439        assert!(family.supports_feature("tool_use"));
440        assert!(!family.supports_feature("unknown"));
441    }
442}