Cli list models

crates/cli/src/interactive/app.rs

@@ -0,0 +1,92 @@
+use std::io::IsTerminal;
+
+use iocraft::prelude::*;
+use uzu::{
+    engine::{Engine, EngineConfig, EngineError},
+    settings::SettingsError,
+};
+
+use crate::interactive::{
+    components::{AppSettings, Application, Preferences, Theme},
+    model::resolve_model_id,
+};
+
+#[derive(Debug, Clone, PartialEq, thiserror::Error)]
+#[non_exhaustive]
+pub enum CliError {
+    #[error(transparent)]
+    Engine(#[from] EngineError),
+    #[error(transparent)]
+    Settigs(#[from] SettingsError),
+    #[error("Rendering error: {message}")]
+    RenderingError {
+        message: String,
+    },
+}
+
+#[derive(Clone)]
+pub struct CliApplication {
+    engine: Engine,
+}
+
+impl CliApplication {
+    pub async fn create(config: EngineConfig) -> Result<Self, CliError> {
+        let engine = Engine::new(config).await?;
+        Ok(Self::new(engine))
+    }
+
+    pub fn new(engine: Engine) -> Self {
+        Self {
+            engine,
+        }
+    }
+
+    pub async fn run_with_model(
+        &self,
+        model: Option<String>,
+    ) -> Result<(), CliError> {
+        if !std::io::stdout().is_terminal() {
+            return Err(CliError::RenderingError {
+                message: "stdout is not a terminal".to_string(),
+            });
+        }
+
+        let settings = self.engine.settings().await.ok();
+        let theme = match &settings {
+            Some(settings) => Theme::load(settings)?.unwrap_or_default(),
+            None => Theme::default(),
+        };
+        let preferences = match &settings {
+            Some(settings) => Preferences::load(settings)?,
+            None => Preferences::default(),
+        };
+        let app_settings = match &settings {
+            Some(settings) => AppSettings::load(settings)?,
+            None => AppSettings::default(),
+        };
+
+        let requested_model = model.or_else(|| app_settings.selected_model_id.clone());
+        let selected_model = match requested_model {
+            Some(model) => resolve_model_id(&self.engine, model).await?,
+            None => None,
+        };
+
+        element! {
+            Application(
+                engine: Some(self.engine.clone()),
+                settings: settings,
+                theme: Some(theme),
+                preferences: Some(preferences),
+                app_settings: Some(app_settings),
+                model: selected_model,
+            )
+        }
+        .render_loop()
+        .await
+        .map_err(|error| CliError::RenderingError {
+            message: error.to_string(),
+        })?;
+
+        Ok(())
+    }
+}

crates/cli/src/interactive/list.rs

@@ -0,0 +1,59 @@
+use std::collections::HashSet;
+
+use shoji::types::model::Model;
+
+pub struct ModelCheckpoint {
+    pub id: String,
+    pub name: String,
+}
+
+pub fn get_checkpoints(
+    models: &[Model],
+    model_id: &str,
+) -> Vec<ModelCheckpoint> {
+    models
+        .iter()
+        .filter(|model| {
+            let Some(family) = &model.family else {
+                return false;
+            };
+            let Some(properties) = &model.properties else {
+                return false;
+            };
+
+            let family_id = family.identifier.rsplit(':').next().unwrap_or(&family.identifier);
+            format!("{family_id}:{}", properties.identifier) == model_id
+        })
+        .map(|model| ModelCheckpoint {
+            id: model.identifier.clone(),
+            name: model.name(),
+        })
+        .collect()
+}
+
+pub struct ModelFamily {
+    pub id: String,
+    pub name: String,
+}
+
+pub fn get_families(models: &[Model]) -> Vec<ModelFamily> {
+    let mut families = Vec::<ModelFamily>::new();
+    let mut ids_set = HashSet::<String>::new();
+
+    for model in models.iter() {
+        if let Some(ref family) = model.family {
+            if let Some(ref properties) = model.properties {
+                let model_id = format!("{}:{}", family.identifier.split(":").last().unwrap(), properties.identifier);
+                if !ids_set.contains(&model_id) {
+                    ids_set.insert(model_id.clone());
+                    families.push(ModelFamily {
+                        id: model_id,
+                        name: format!("{} {}", family.metadata.name, properties.metadata.name),
+                    })
+                }
+            }
+        }
+    }
+
+    families
+}

crates/cli/src/interactive/mod.rs

@@ -1,103 +1,76 @@
+use comfy_table::{
+    ContentArrangement, Table,
+    modifiers::{UTF8_ROUND_CORNERS, UTF8_SOLID_INNER_BORDERS},
+    presets::UTF8_FULL,
+};
+use uzu::engine::{Engine, EngineConfig};
+
+use crate::interactive::{
+    app::CliApplication,
+    list::{get_checkpoints, get_families},
+};
+
+mod app;
 mod components;
 mod flows;
 mod helpers;
+mod list;
+mod model;
 mod sessions;
 
-use std::io::IsTerminal;
-
-use components::{Application, Preferences, Theme};
-use iocraft::prelude::*;
-use uzu::{
-    engine::{Engine, EngineConfig, EngineError},
-    settings::SettingsError,
-};
-
-use crate::interactive::components::AppSettings;
-
-#[derive(Debug, Clone, PartialEq, thiserror::Error)]
-#[non_exhaustive]
-pub enum CliError {
-    #[error(transparent)]
-    Engine(#[from] EngineError),
-    #[error(transparent)]
-    Settigs(#[from] SettingsError),
-    #[error("Rendering error: {message}")]
-    RenderingError {
-        message: String,
-    },
-}
-
-#[derive(Clone)]
-pub struct CliApplication {
-    engine: Engine,
+pub async fn run_interactive(model: Option<String>) -> anyhow::Result<()> {
+    let engine_config = EngineConfig::default().with_application_identifier("com.trymirai.cli".to_string());
+    let application = CliApplication::create(engine_config).await?;
+    application.run_with_model(model).await?;
+    Ok(())
 }
 
-impl CliApplication {
-    pub fn new(engine: Engine) -> Self {
-        Self {
-            engine,
-        }
+pub async fn run_list_models() -> anyhow::Result<()> {
+    let engine_config = EngineConfig::default().with_application_identifier("com.trymirai.cli".to_string());
+    let engine = Engine::new(engine_config).await?;
+    let models = engine.models().await?;
+    if models.is_empty() {
+        return Err(anyhow::anyhow!("No models to run"));
     }
 
-    pub async fn run_with_model(
-        &self,
-        model: Option<String>,
-    ) -> Result<(), CliError> {
-        if !std::io::stdout().is_terminal() {
-            return Err(CliError::RenderingError {
-                message: "stdout is not a terminal".to_string(),
-            });
-        }
+    let families = get_families(&models);
+    let mut table = Table::new();
+    table
+        .load_preset(UTF8_FULL)
+        .apply_modifier(UTF8_ROUND_CORNERS)
+        .apply_modifier(UTF8_SOLID_INNER_BORDERS)
+        .set_content_arrangement(ContentArrangement::Dynamic)
+        .set_header(vec!["Name", "ID"]);
 
-        let settings = self.engine.settings().await.ok();
-        let theme = match &settings {
-            Some(settings) => Theme::load(settings)?.unwrap_or_default(),
-            None => Theme::default(),
-        };
-        let preferences = match &settings {
-            Some(settings) => Preferences::load(settings)?,
-            None => Preferences::default(),
-        };
-        let app_settings = match &settings {
-            Some(settings) => AppSettings::load(settings)?,
-            None => AppSettings::default(),
-        };
-
-        let mut selected_model = model;
-        if selected_model.is_none() {
-            selected_model = app_settings.selected_model_id.clone();
-        }
+    for family in &families {
+        table.add_row(vec![&family.name, &family.id]);
+    }
+    println!("{table}");
 
-        element! {
-            Application(
-                engine: Some(self.engine.clone()),
-                settings: settings,
-                theme: Some(theme),
-                preferences: Some(preferences),
-                app_settings: Some(app_settings),
-                model: selected_model,
-            )
-        }
-        .render_loop()
-        .await
-        .map_err(|error| CliError::RenderingError {
-            message: error.to_string(),
-        })?;
+    Ok(())
+}
 
-        Ok(())
+pub async fn run_list_checkpoints(model_id: String) -> anyhow::Result<()> {
+    let engine_config = EngineConfig::default().with_application_identifier("com.trymirai.cli".to_string());
+    let engine = Engine::new(engine_config).await?;
+    let models = engine.models().await?;
+    let checkpoints = get_checkpoints(&models, &model_id);
+    if checkpoints.is_empty() {
+        return Err(anyhow::anyhow!("No checkpoints found for model: {model_id}"));
     }
-}
 
-impl CliApplication {
-    pub async fn create(config: EngineConfig) -> Result<Self, CliError> {
-        let engine = Engine::new(config).await?;
-        Ok(Self::new(engine))
+    let mut table = Table::new();
+    table
+        .load_preset(UTF8_FULL)
+        .apply_modifier(UTF8_ROUND_CORNERS)
+        .apply_modifier(UTF8_SOLID_INNER_BORDERS)
+        .set_content_arrangement(ContentArrangement::Dynamic)
+        .set_header(vec!["Name", "ID"]);
+
+    for checkpoint in &checkpoints {
+        table.add_row(vec![&checkpoint.name, &checkpoint.id]);
     }
-}
+    println!("{table}");
 
-pub async fn run_interactive(model: Option<String>) -> anyhow::Result<()> {
-    let engine_config = EngineConfig::default().with_application_identifier("com.trymirai.cli".to_string());
-    let application = CliApplication::create(engine_config).await?;
-    application.run_with_model(model).await?;
     Ok(())
 }