agentic_core/
models.rs

1use anyhow::Result;
2use reqwest::Client;
3use serde::{Deserialize, Serialize};
4use serde_json::Value;
5use std::time::Duration;
6
7#[derive(Debug, Clone, Serialize, Deserialize)]
8pub struct AtomicNote {
9    pub header_tags: Vec<String>,
10    pub body_text: String,
11}
12
13#[derive(Debug, Clone, Serialize, Deserialize)]
14pub struct OllamaModel {
15    pub name: String,
16    pub size: String,
17    pub modified: String,
18}
19
20#[derive(Debug, Clone, Serialize, Deserialize)]
21pub struct OpenRouterModel {
22    pub id: String,
23    pub name: String,
24    pub description: String,
25    pub pricing: ModelPricing,
26    pub context_length: u32,
27}
28
29#[derive(Debug, Clone, Serialize, Deserialize)]
30pub struct ModelPricing {
31    pub prompt: String,
32    pub completion: String,
33}
34
35#[derive(Debug, Serialize, Deserialize)]
36struct OllamaListResponse {
37    models: Vec<OllamaModelRaw>,
38}
39
40#[derive(Debug, Serialize, Deserialize)]
41struct OllamaModelRaw {
42    name: String,
43    size: i64,
44    modified_at: String,
45}
46
47#[derive(Debug, Serialize, Deserialize)]
48struct OpenRouterListResponse {
49    data: Vec<OpenRouterModelRaw>,
50}
51
52#[derive(Debug, Serialize, Deserialize)]
53struct OpenRouterModelRaw {
54    id: String,
55    name: String,
56    description: Option<String>,
57    pricing: ModelPricingRaw,
58    context_length: u32,
59}
60
61#[derive(Debug, Serialize, Deserialize)]
62struct ModelPricingRaw {
63    prompt: String,
64    completion: String,
65}
66
67pub struct ModelValidator {
68    client: Client,
69}
70
71impl ModelValidator {
72    pub fn new() -> Self {
73        let client = Client::builder()
74            .timeout(Duration::from_secs(5))
75            .build()
76            .unwrap_or_default();
77
78        Self { client }
79    }
80
81    pub async fn fetch_ollama_models(&self, endpoint: &str) -> Result<Vec<OllamaModel>> {
82        let url = if endpoint.starts_with("http") {
83            format!("{}/api/tags", endpoint)
84        } else {
85            format!("http://{}/api/tags", endpoint)
86        };
87
88        let response = self.client.get(&url).send().await?;
89
90        if !response.status().is_success() {
91            return Err(anyhow::anyhow!("Ollama endpoint not accessible"));
92        }
93
94        let ollama_response: OllamaListResponse = response.json().await?;
95
96        let models = ollama_response
97            .models
98            .into_iter()
99            .map(|raw| OllamaModel {
100                name: raw.name,
101                size: format_size(raw.size),
102                modified: format_relative_time(&raw.modified_at),
103            })
104            .collect();
105
106        Ok(models)
107    }
108
109    pub async fn fetch_openrouter_models(&self, api_key: &str) -> Result<Vec<OpenRouterModel>> {
110        let url = "https://openrouter.ai/api/v1/models";
111
112        let response = self
113            .client
114            .get(url)
115            .header("Authorization", format!("Bearer {}", api_key))
116            .send()
117            .await?;
118
119        if !response.status().is_success() {
120            return Err(anyhow::anyhow!(
121                "OpenRouter API not accessible or invalid API key"
122            ));
123        }
124
125        let openrouter_response: OpenRouterListResponse = response.json().await?;
126
127        let mut models: Vec<OpenRouterModel> = openrouter_response
128            .data
129            .into_iter()
130            .map(|raw| OpenRouterModel {
131                id: raw.id,
132                name: raw.name,
133                description: raw
134                    .description
135                    .unwrap_or_else(|| "No description available".to_string()),
136                pricing: ModelPricing {
137                    prompt: raw.pricing.prompt,
138                    completion: raw.pricing.completion,
139                },
140                context_length: raw.context_length,
141            })
142            .collect();
143
144        // Sort models: free models first, then paid models
145        models.sort_by(|a, b| {
146            let a_is_free = a.pricing.prompt == "0" && a.pricing.completion == "0";
147            let b_is_free = b.pricing.prompt == "0" && b.pricing.completion == "0";
148
149            match (a_is_free, b_is_free) {
150                (true, false) => std::cmp::Ordering::Less, // Free comes first
151                (false, true) => std::cmp::Ordering::Greater, // Paid comes after
152                _ => a.name.cmp(&b.name),                  // Same type, sort by name
153            }
154        });
155
156        Ok(models)
157    }
158
159    pub async fn validate_local_endpoint(&self, endpoint: &str, model: &str) -> Result<()> {
160        let url = if endpoint.starts_with("http") {
161            format!("{}/api/tags", endpoint)
162        } else {
163            format!("http://{}/api/tags", endpoint)
164        };
165
166        let response = self.client.get(&url).send().await?;
167
168        if !response.status().is_success() {
169            return Err(anyhow::anyhow!("Local endpoint not accessible"));
170        }
171
172        let models: Value = response.json().await?;
173
174        if let Some(models_array) = models.get("models").and_then(|m| m.as_array()) {
175            let model_exists = models_array.iter().any(|m| {
176                m.get("name")
177                    .and_then(|name| name.as_str())
178                    .map(|name| name == model)
179                    .unwrap_or(false)
180            });
181
182            if model_exists {
183                Ok(())
184            } else {
185                Err(anyhow::anyhow!(
186                    "Model '{}' not found on local endpoint",
187                    model
188                ))
189            }
190        } else {
191            Err(anyhow::anyhow!(
192                "Invalid response format from local endpoint"
193            ))
194        }
195    }
196
197    pub async fn validate_cloud_endpoint(&self, api_key: &str, model: &str) -> Result<()> {
198        let url = "https://openrouter.ai/api/v1/models";
199
200        let response = self
201            .client
202            .get(url)
203            .header("Authorization", format!("Bearer {}", api_key))
204            .send()
205            .await?;
206
207        if !response.status().is_success() {
208            return Err(anyhow::anyhow!(
209                "Cloud API key invalid or endpoint not accessible"
210            ));
211        }
212
213        let models: Value = response.json().await?;
214
215        if let Some(models_array) = models.get("data").and_then(|m| m.as_array()) {
216            let model_exists = models_array.iter().any(|m| {
217                m.get("id")
218                    .and_then(|id| id.as_str())
219                    .map(|id| id == model)
220                    .unwrap_or(false)
221            });
222
223            if model_exists {
224                Ok(())
225            } else {
226                Err(anyhow::anyhow!("Model '{}' not found in OpenRouter", model))
227            }
228        } else {
229            Err(anyhow::anyhow!("Invalid response format from OpenRouter"))
230        }
231    }
232
233    pub async fn test_local_generation(&self, endpoint: &str, model: &str) -> Result<()> {
234        let url = if endpoint.starts_with("http") {
235            format!("{}/api/generate", endpoint)
236        } else {
237            format!("http://{}/api/generate", endpoint)
238        };
239
240        let payload = serde_json::json!({
241            "model": model,
242            "prompt": "Hello",
243            "stream": false,
244            "options": {
245                "num_predict": 1
246            }
247        });
248
249        let response = self.client.post(&url).json(&payload).send().await?;
250
251        if response.status().is_success() {
252            Ok(())
253        } else {
254            Err(anyhow::anyhow!(
255                "Failed to generate response from local model"
256            ))
257        }
258    }
259
260    pub async fn test_cloud_generation(&self, api_key: &str, model: &str) -> Result<()> {
261        let url = "https://openrouter.ai/api/v1/chat/completions";
262
263        let payload = serde_json::json!({
264            "model": model,
265            "messages": [{"role": "user", "content": "Hello"}],
266            "max_tokens": 1
267        });
268
269        let response = self
270            .client
271            .post(url)
272            .header("Authorization", format!("Bearer {}", api_key))
273            .header("Content-Type", "application/json")
274            .json(&payload)
275            .send()
276            .await?;
277
278        if response.status().is_success() {
279            Ok(())
280        } else {
281            Err(anyhow::anyhow!(
282                "Failed to generate response from cloud model"
283            ))
284        }
285    }
286}
287
288#[derive(Serialize)]
289struct LocalGenerationRequest<'a> {
290    model: &'a str,
291    prompt: &'a str,
292    stream: bool,
293}
294
295#[derive(Deserialize)]
296struct LocalGenerationResponse {
297    response: String,
298}
299
300pub async fn call_local_model(
301    endpoint: &str,
302    model: &str,
303    prompt: &str,
304) -> Result<String, anyhow::Error> {
305    let client = Client::new();
306    let url = if endpoint.starts_with("http") {
307        format!("{}/api/generate", endpoint)
308    } else {
309        format!("http://{}/api/generate", endpoint)
310    };
311
312    let payload = LocalGenerationRequest {
313        model,
314        prompt,
315        stream: false,
316    };
317
318    let response = client.post(&url).json(&payload).send().await?;
319
320    if response.status().is_success() {
321        let gen_response: LocalGenerationResponse = response.json().await?;
322        Ok(gen_response.response)
323    } else {
324        Err(anyhow::anyhow!(
325            "Failed to get response from local model. Status: {}",
326            response.status()
327        ))
328    }
329}
330
331impl Default for ModelValidator {
332    fn default() -> Self {
333        Self::new()
334    }
335}
336
337fn format_size(bytes: i64) -> String {
338    const UNITS: &[&str] = &["B", "KB", "MB", "GB", "TB"];
339    let mut size = bytes as f64;
340    let mut unit_index = 0;
341
342    while size >= 1024.0 && unit_index < UNITS.len() - 1 {
343        size /= 1024.0;
344        unit_index += 1;
345    }
346
347    if unit_index == 0 {
348        format!("{} {}", size as i64, UNITS[unit_index])
349    } else {
350        format!("{:.1} {}", size, UNITS[unit_index])
351    }
352}
353
354fn format_relative_time(_iso_time: &str) -> String {
355    // For now, just return a simple format
356    // Parse ISO time and return relative time
357    "recently".to_string()
358}
agentic_core/models.rs

agentic_core/
models.rs
