kowalski_core/model/
mod.rs

1use crate::error::KowalskiError;
2use reqwest::Client;
3use serde::{Deserialize, Serialize};
4
5#[derive(Debug, Serialize, Deserialize)]
6pub struct ModelInfo {
7    pub name: String,
8    pub size: u64,
9    pub digest: String,
10    pub modified_at: String,
11}
12
13#[derive(Debug, Serialize, Deserialize)]
14pub struct ModelsResponse {
15    pub models: Vec<ModelInfo>,
16}
17
18#[derive(Debug, Serialize, Deserialize)]
19pub struct PullResponse {
20    pub status: String,
21    pub digest: Option<String>,
22    pub total: Option<u64>,
23    pub completed: Option<u64>,
24}
25
26/// The main struct that makes our AI models feel special.
27pub struct ModelManager {
28    client: Client,
29    base_url: String,
30}
31
32impl ModelManager {
33    /// Creates a new model manager with the specified base URL.
34    pub fn new(base_url: String) -> Result<Self, KowalskiError> {
35        let client = reqwest::ClientBuilder::new()
36            .pool_max_idle_per_host(0)
37            .build()
38            .map_err(KowalskiError::Request)?;
39
40        Ok(Self { client, base_url })
41    }
42
43    /// Lists available models
44    pub async fn list_models(&self) -> Result<ModelsResponse, KowalskiError> {
45        let response = self
46            .client
47            .get(format!("{}/api/tags", self.base_url))
48            .send()
49            .await?;
50
51        if !response.status().is_success() {
52            let error_text = response.text().await?;
53            return Err(KowalskiError::Server(error_text));
54        }
55
56        let models_response: ModelsResponse = response.json().await?;
57        Ok(models_response)
58    }
59
60    /// Checks if a model exists
61    pub async fn model_exists(&self, model_name: &str) -> Result<bool, KowalskiError> {
62        let models = self.list_models().await?;
63        Ok(models.models.iter().any(|m| m.name == model_name))
64    }
65
66    /// Pulls a model from the server
67    pub async fn pull_model(&self, model_name: &str) -> Result<PullResponse, KowalskiError> {
68        let response = self
69            .client
70            .post(format!("{}/api/pull", self.base_url))
71            .json(&serde_json::json!({
72                "name": model_name
73            }))
74            .send()
75            .await?;
76
77        if !response.status().is_success() {
78            let error_text = response.text().await?;
79            return Err(KowalskiError::Server(error_text));
80        }
81
82        let pull_response: PullResponse = response.json().await?;
83        Ok(pull_response)
84    }
85}
kowalski_core/model/mod.rs

kowalski_core/model/
mod.rs
