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    /// Optional legacy token threshold for automatic compaction. Runtime
67    /// compaction uses the resolved provider/session budget instead of
68    /// deriving a percentage from `context_window`.
69    auto_compact_token_limit: Option<i64>,
70
71    /// Whether the model supports reasoning summaries
72    pub supports_reasoning_summaries: bool,
73
74    /// Default reasoning effort for this model family
75    default_reasoning_effort: Option<ReasoningEffortLevel>,
76
77    /// Whether the model supports parallel tool calls
78    supports_parallel_tool_calls: bool,
79
80    /// Whether the model needs special apply_patch instructions
81    needs_special_apply_patch_instructions: bool,
82
83    /// Preferred shell tool type for this model family
84    shell_type: ShellToolType,
85
86    /// Truncation policy for model output
87    truncation_policy: TruncationPolicy,
88
89    /// Names of experimental tools supported by this model family
90    experimental_supported_tools: Vec<String>,
91
92    /// Percentage of context window considered usable for inputs
93    effective_context_window_percent: i64,
94
95    /// Whether the model supports verbosity settings
96    support_verbosity: bool,
97
98    /// Whether the model supports tool use
99    supports_tool_use: bool,
100
101    /// Whether the model supports streaming
102    supports_streaming: bool,
103
104    /// Whether the model supports thinking/reasoning output
105    supports_thinking: bool,
106}
107
108impl Default for ModelFamily {
109    fn default() -> Self {
110        Self {
111            slug: String::new(),
112            family: String::new(),
113            provider: Provider::default(),
114            context_window: Some(DEFAULT_CONTEXT_WINDOW),
115            auto_compact_token_limit: None,
116            supports_reasoning_summaries: false,
117            default_reasoning_effort: None,
118            supports_parallel_tool_calls: false,
119            needs_special_apply_patch_instructions: false,
120            shell_type: ShellToolType::Default,
121            truncation_policy: TruncationPolicy::default(),
122            experimental_supported_tools: Vec::new(),
123            effective_context_window_percent: 95,
124            support_verbosity: false,
125            supports_tool_use: true,
126            supports_streaming: true,
127            supports_thinking: false,
128        }
129    }
130}
131
132impl ModelFamily {
133    /// Create a new model family with the given slug
134    fn new(slug: impl Into<String>, family: impl Into<String>, provider: Provider) -> Self {
135        Self {
136            slug: slug.into(),
137            family: family.into(),
138            provider,
139            ..Default::default()
140        }
141    }
142
143    /// Get the explicitly configured legacy auto-compact token limit.
144    fn auto_compact_token_limit(&self) -> Option<i64> {
145        self.auto_compact_token_limit
146    }
147
148    /// Get the model slug
149    pub fn get_model_slug(&self) -> &str {
150        &self.slug
151    }
152
153    /// Check if this family supports a specific feature
154    fn supports_feature(&self, feature: &str) -> bool {
155        match feature {
156            "reasoning" | "thinking" => self.supports_thinking,
157            "tool_use" | "tools" => self.supports_tool_use,
158            "streaming" => self.supports_streaming,
159            "parallel_tools" => self.supports_parallel_tool_calls,
160            _ => self.experimental_supported_tools.contains(&feature.to_string()),
161        }
162    }
163}
164
165/// Macro to simplify model family definitions
166#[macro_export]
167macro_rules! model_family {
168    (
169        $slug:expr, $family:expr, $provider:expr $(, $key:ident : $value:expr )* $(,)?
170    ) => {{
171        let mut mf = $crate::model_family::ModelFamily::new($slug, $family, $provider);
172        $(
173            mf.$key = $value;
174        )*
175        mf
176    }};
177}
178
179/// Internal helper that returns a `ModelFamily` for the given model slug.
180pub fn find_family_for_model(slug: &str) -> ModelFamily {
181    if let Some((provider, raw_slug)) = opencode_provider_and_raw_slug(slug) {
182        let mut family = find_family_for_model(raw_slug);
183        family.slug = slug.to_string();
184        family.provider = provider;
185        return family;
186    }
187
188    // Gemini models
189    if slug.starts_with("gemini-3") {
190        return model_family!(
191            slug, "gemini-3", Provider::Gemini,
192            context_window: Some(LARGE_CONTEXT_WINDOW),
193            supports_thinking: true,
194            supports_parallel_tool_calls: true,
195            supports_reasoning_summaries: true,
196        );
197    }
198    if slug.starts_with("gemini") {
199        return model_family!(
200            slug, "gemini", Provider::Gemini,
201            context_window: Some(LARGE_CONTEXT_WINDOW),
202        );
203    }
204
205    // OpenAI models
206    if slug.starts_with("gpt-6") {
207        return model_family!(
208            slug, "gpt-6", Provider::OpenAI,
209            context_window: Some(LARGE_CONTEXT_WINDOW),
210            supports_thinking: true,
211            supports_parallel_tool_calls: true,
212        );
213    }
214    if slug.starts_with("gpt-5") {
215        return model_family!(
216            slug, "gpt-5", Provider::OpenAI,
217            context_window: Some(DEFAULT_CONTEXT_WINDOW),
218            supports_thinking: true,
219            supports_parallel_tool_calls: true,
220        );
221    }
222    if slug.starts_with("codex") {
223        return model_family!(
224            slug, "codex", Provider::OpenAI,
225            context_window: Some(MEDIUM_CONTEXT_WINDOW),
226            supports_thinking: true,
227            shell_type: ShellToolType::ExecCommand,
228        );
229    }
230    if slug.starts_with("gpt-oss") || slug.contains("gpt-oss") {
231        return model_family!(
232            slug, "gpt-oss", Provider::OpenAI,
233            context_window: Some(96_000),
234        );
235    }
236    if slug.starts_with("o3") || slug.starts_with("o4") {
237        return model_family!(
238            slug, "o-series", Provider::OpenAI,
239            context_window: Some(MEDIUM_CONTEXT_WINDOW),
240            supports_thinking: true,
241            supports_reasoning_summaries: true,
242            needs_special_apply_patch_instructions: true,
243        );
244    }
245
246    // Anthropic models
247    if slug.starts_with("claude-opus") || slug.contains("opus") {
248        return model_family!(
249            slug, "claude-opus", Provider::Anthropic,
250            context_window: Some(MEDIUM_CONTEXT_WINDOW),
251            supports_thinking: true,
252            supports_parallel_tool_calls: true,
253        );
254    }
255    if slug.starts_with("claude-sonnet") || slug.contains("sonnet") {
256        return model_family!(
257            slug, "claude-sonnet", Provider::Anthropic,
258            context_window: Some(MEDIUM_CONTEXT_WINDOW),
259            supports_thinking: true,
260        );
261    }
262    if slug.starts_with("claude-haiku") || slug.contains("haiku") {
263        return model_family!(
264            slug, "claude-haiku", Provider::Anthropic,
265            context_window: Some(MEDIUM_CONTEXT_WINDOW),
266        );
267    }
268    if slug.starts_with("claude") {
269        return model_family!(
270            slug, "claude", Provider::Anthropic,
271            context_window: Some(MEDIUM_CONTEXT_WINDOW),
272        );
273    }
274
275    // DeepSeek models
276    if slug.contains("deepseek") && slug.contains("reason") {
277        return model_family!(
278            slug, "deepseek-reasoner", Provider::DeepSeek,
279            context_window: Some(DEFAULT_CONTEXT_WINDOW),
280            supports_thinking: true,
281        );
282    }
283    if slug.contains("deepseek") {
284        return model_family!(
285            slug, "deepseek", Provider::DeepSeek,
286            context_window: Some(DEFAULT_CONTEXT_WINDOW),
287        );
288    }
289
290    // Official Meta AI Muse models
291    if slug.starts_with("muse-spark-") {
292        return model_family!(
293            slug, "muse-spark", Provider::Meta,
294            context_window: Some(LARGE_CONTEXT_WINDOW),
295            supports_thinking: true,
296            supports_parallel_tool_calls: true,
297            supports_reasoning_summaries: false,
298        );
299    }
300
301    // Z.AI GLM models
302    if slug.contains("glm-5") {
303        return model_family!(
304            slug, "glm-5", Provider::ZAI,
305            context_window: Some(DEFAULT_CONTEXT_WINDOW),
306            supports_thinking: true,
307        );
308    }
309    if slug.contains("glm") {
310        return model_family!(
311            slug, "glm", Provider::ZAI,
312            context_window: Some(DEFAULT_CONTEXT_WINDOW),
313        );
314    }
315
316    // MiniMax models
317    if slug.contains("minimax") {
318        return model_family!(
319            slug, "minimax", Provider::Minimax,
320            context_window: Some(DEFAULT_CONTEXT_WINDOW),
321            supports_thinking: true,
322        );
323    }
324
325    // Moonshot/Kimi models
326    if slug.contains("kimi") || slug.contains("moonshot") {
327        return model_family!(
328            slug, "kimi", Provider::Moonshot,
329            context_window: Some(DEFAULT_CONTEXT_WINDOW),
330            supports_thinking: slug.contains("thinking"),
331        );
332    }
333
334    // Qwen models (via OpenRouter or Ollama)
335    if slug.contains("qwen") {
336        return model_family!(
337            slug, "qwen", Provider::OpenRouter,
338            context_window: Some(DEFAULT_CONTEXT_WINDOW),
339            supports_thinking: slug.contains("thinking"),
340        );
341    }
342
343    // Ollama local models
344    if slug.starts_with("ollama/") || slug.contains(":") {
345        return model_family!(
346            slug, "ollama-local", Provider::Ollama,
347            context_window: Some(DEFAULT_CONTEXT_WINDOW),
348        );
349    }
350
351    // OpenRouter models (fallback for unrecognized patterns)
352    if slug.contains("/") {
353        return model_family!(
354            slug, "openrouter", Provider::OpenRouter,
355            context_window: Some(DEFAULT_CONTEXT_WINDOW),
356        );
357    }
358
359    // Default fallback
360    model_family!(
361        slug, "unknown", Provider::default(),
362        context_window: Some(DEFAULT_CONTEXT_WINDOW),
363    )
364}
365
366fn opencode_provider_and_raw_slug(slug: &str) -> Option<(Provider, &str)> {
367    if let Some(raw_slug) = slug.strip_prefix("opencode-go/") {
368        Some((Provider::OpenCodeGo, raw_slug))
369    } else if let Some(raw_slug) = slug.strip_prefix("opencode/").or_else(|| slug.strip_prefix("opencode-zen/")) {
370        Some((Provider::OpenCodeZen, raw_slug))
371    } else {
372        None
373    }
374}
375
376#[cfg(test)]
377mod tests {
378    use super::*;
379
380    #[test]
381    fn test_gemini_family_detection() {
382        let family = find_family_for_model("gemini-3-flash-preview");
383        assert_eq!(family.family, "gemini-3");
384        assert_eq!(family.provider, Provider::Gemini);
385        assert!(family.context_window.unwrap() >= LARGE_CONTEXT_WINDOW);
386    }
387
388    #[test]
389    fn test_gpt5_family_detection() {
390        let family = find_family_for_model("gpt-5-codex");
391        assert_eq!(family.family, "gpt-5");
392        assert_eq!(family.provider, Provider::OpenAI);
393        assert!(family.supports_thinking);
394    }
395
396    #[test]
397    fn test_claude_family_detection() {
398        let family = find_family_for_model("claude-opus-4.5");
399        assert_eq!(family.family, "claude-opus");
400        assert_eq!(family.provider, Provider::Anthropic);
401    }
402
403    #[test]
404    fn test_opencode_zen_family_detection_preserves_provider() {
405        let family = find_family_for_model("opencode/gpt-5.4");
406        assert_eq!(family.family, "gpt-5");
407        assert_eq!(family.provider, Provider::OpenCodeZen);
408        assert!(family.supports_thinking);
409    }
410
411    #[test]
412    fn test_opencode_go_family_detection_preserves_provider() {
413        let family = find_family_for_model("opencode-go/kimi-k2.5");
414        assert_eq!(family.family, "kimi");
415        assert_eq!(family.provider, Provider::OpenCodeGo);
416    }
417
418    #[test]
419    fn test_auto_compact_limit_requires_explicit_value() {
420        let family = ModelFamily {
421            context_window: Some(100_000),
422            ..Default::default()
423        };
424        assert_eq!(family.auto_compact_token_limit(), None);
425
426        let family = ModelFamily {
427            context_window: Some(100_000),
428            auto_compact_token_limit: Some(90_000),
429            ..Default::default()
430        };
431        assert_eq!(family.auto_compact_token_limit(), Some(90_000));
432    }
433
434    #[test]
435    fn test_supports_feature() {
436        let family = ModelFamily {
437            supports_thinking: true,
438            supports_tool_use: true,
439            ..Default::default()
440        };
441        assert!(family.supports_feature("thinking"));
442        assert!(family.supports_feature("tool_use"));
443        assert!(!family.supports_feature("unknown"));
444    }
445}