@sigit / sigit
Code Commits Pulls CI/CD
main
development feature/crates-release feature/v1.0.4 main
rs 168 lines 6.38 KB
Raw
1 //! Model picker types and item construction, shared across platforms.
2 //! The unix-only TUI in `chat.rs` pulls from here.
3
4 use onde::inference::GgufModelConfig;
5
6 use crate::setup::DiscoveredModel;
7
8 pub(crate) use crate::setup::ModelCacheHealth;
9
10 // ── Types ─────────────────────────────────────────────────────────────────────
11
12 #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
13 pub(crate) enum ModelSource {
14 Onde,
15 HuggingFace,
16 /// not downloaded yet — selecting it triggers a download into the app-group cache.
17 Available,
18 Fallback,
19 }
20
21 #[derive(Clone)]
22 pub(crate) struct ModelPickerItem {
23 pub(crate) display_name: String,
24 pub(crate) description: String,
25 pub(crate) tool_calling: bool,
26 pub(crate) max_tokens: u64,
27 pub(crate) config: GgufModelConfig,
28 pub(crate) source_label: String,
29
30 pub(crate) source: ModelSource,
31 pub(crate) cache_health: ModelCacheHealth,
32 }
33
34 // ── Model ID → GgufModelConfig mapping ────────────────────────────────────────
35
36 /// map a HF model ID to its config constructor, or `None` if we don't support it.
37 pub(crate) fn model_id_to_config(model_id: &str) -> Option<GgufModelConfig> {
38 Some(match model_id {
39 "bartowski/Qwen_Qwen3-4B-GGUF" => GgufModelConfig::qwen3_4b(),
40 "bartowski/Qwen_Qwen3-8B-GGUF" => GgufModelConfig::qwen3_8b(),
41 "bartowski/Qwen_Qwen3-14B-GGUF" => GgufModelConfig::qwen3_14b(),
42 "bartowski/Qwen_Qwen3-1.7B-GGUF" => GgufModelConfig::qwen3_1_7b(),
43 "bartowski/Qwen2.5-3B-Instruct-GGUF" => GgufModelConfig::qwen25_3b(),
44 "bartowski/Qwen2.5-1.5B-Instruct-GGUF" => GgufModelConfig::qwen25_1_5b(),
45 "bartowski/Qwen2.5-Coder-3B-Instruct-GGUF" => GgufModelConfig::qwen25_coder_3b(),
46 "bartowski/Qwen2.5-Coder-1.5B-Instruct-GGUF" => GgufModelConfig::qwen25_coder_1_5b(),
47 "bartowski/Qwen2.5-Coder-7B-Instruct-GGUF" => GgufModelConfig::qwen25_coder_7b(),
48 "TheBloke/deepseek-coder-6.7B-instruct-GGUF" => GgufModelConfig::deepseek_coder_6_7b(),
49 _ => return None,
50 })
51 }
52
53 fn is_tool_calling(model_id: &str) -> bool {
54 matches!(
55 model_id,
56 "bartowski/Qwen_Qwen3-4B-GGUF"
57 | "bartowski/Qwen_Qwen3-8B-GGUF"
58 | "bartowski/Qwen_Qwen3-14B-GGUF"
59 | "bartowski/Qwen_Qwen3-1.7B-GGUF"
60 | "bartowski/Qwen2.5-Coder-7B-Instruct-GGUF"
61 )
62 }
63
64 /// tool-calling models get more tokens because `<think>` blocks eat into the budget.
65 fn max_tokens_for(model_id: &str) -> u64 {
66 if is_tool_calling(model_id) { 4096 } else { 512 }
67 }
68
69 // ── Builder ───────────────────────────────────────────────────────────────────
70
71 /// collect every model the picker should show: local cache, remote available, fallback.
72 /// sorted by source (Onde > HF > Available > Fallback), then alphabetically.
73 pub(crate) fn build_model_picker_items() -> Vec<ModelPickerItem> {
74 let mut items = Vec::new();
75
76 // ── 1. Locally discovered models ─────────────────────────────────────
77 for discovered in crate::setup::discover_local_models() {
78 if let Some(item) = discovered_model_to_picker_item(discovered) {
79 items.push(item);
80 }
81 }
82
83 // ── 2. Supported models not yet downloaded ───────────────────────────
84 for info in onde::inference::models::SUPPORTED_MODEL_INFO {
85 let already_present = items.iter().any(|item| item.config.model_id == info.id);
86 if already_present {
87 continue;
88 }
89
90 let config = match model_id_to_config(info.id) {
91 Some(config) => config,
92 None => continue,
93 };
94
95 let tool_calling = is_tool_calling(info.id);
96 let max_tokens = max_tokens_for(info.id);
97
98 items.push(ModelPickerItem {
99 display_name: config.display_name.clone(),
100 description: config.approx_memory.clone(),
101 tool_calling,
102 max_tokens,
103 config,
104 source_label: "Onde".to_string(),
105
106 source: ModelSource::Available,
107 cache_health: ModelCacheHealth::NotDownloaded,
108 });
109 }
110
111 // ── 3. Fallback ──────────────────────────────────────────────────────
112 if items.is_empty() {
113 let config = GgufModelConfig::platform_default();
114 let tool_calling = is_tool_calling(&config.model_id);
115 let max_tokens = max_tokens_for(&config.model_id);
116
117 items.push(ModelPickerItem {
118 display_name: config.display_name.clone(),
119 description: config.approx_memory.clone(),
120 tool_calling,
121 max_tokens,
122 config,
123 source_label: "Platform default".to_string(),
124
125 source: ModelSource::Fallback,
126 cache_health: ModelCacheHealth::Complete,
127 });
128 }
129
130 items.sort_by(|left, right| {
131 left.source
132 .cmp(&right.source)
133 .then_with(|| left.display_name.cmp(&right.display_name))
134 });
135
136 items
137 }
138
139 // ── Internal helpers ──────────────────────────────────────────────────────────
140
141 fn discovered_model_to_picker_item(model: DiscoveredModel) -> Option<ModelPickerItem> {
142 let source_label = if model.from_app_group {
143 "Onde".to_string()
144 } else {
145 "HuggingFace".to_string()
146 };
147
148 let config = model_id_to_config(&model.model_id)?;
149
150 let tool_calling = is_tool_calling(&model.model_id);
151 let max_tokens = max_tokens_for(&model.model_id);
152
153 Some(ModelPickerItem {
154 display_name: config.display_name.clone(),
155 description: config.approx_memory.clone(),
156 tool_calling,
157 max_tokens,
158 config,
159 source_label,
160
161 source: if model.from_app_group {
162 ModelSource::Onde
163 } else {
164 ModelSource::HuggingFace
165 },
166 cache_health: model.cache_health,
167 })
168 }