diff --git a/Cargo.toml b/Cargo.toml index 00de523..6b59185 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -51,6 +51,7 @@ colored = "2.2.0" ollama-rs = { version = "0.2.5", features = ["rustls", "stream"] } dkn-workflows = { git = "https://github.com/firstbatchxyz/dkn-compute-node" } reqwest = { version = "0.12.12", features = ["rustls-tls", "json"] } +prettytable = "0.10.0" # crypto stuff sha3 = "0.10.8" diff --git a/install.sh b/install.sh index 52e8180..ad59ff7 100755 --- a/install.sh +++ b/install.sh @@ -118,7 +118,7 @@ install_binary() { SHELL_NAME=$(basename "$SHELL") print_step "Detected shell: $SHELL_NAME" - # Check shell and corresponding config files + # check shell and corresponding config files case "$SHELL_NAME" in "fish") CONFIG_FILE="$HOME/.config/fish/config.fish" @@ -127,7 +127,7 @@ install_binary() { CONFIG_FILE="$HOME/.zshrc" ;; "bash") - # Check for both .bash_profile and .bashrc + # check for both .bash_profile and .bashrc if [ -f "$HOME/.bash_profile" ]; then CONFIG_FILE="$HOME/.bash_profile" else @@ -142,7 +142,7 @@ install_binary() { ;; esac - # If config file exists, add PATH if not already present + # if config file exists, add PATH if not already present if [ -f "$CONFIG_FILE" ]; then if grep -q "export PATH=\"${DRIA_INSTALL_DIR}:\$PATH\"" "$CONFIG_FILE"; then print_step "Dria Compute Launcher path exists in $CONFIG_FILE" @@ -153,9 +153,10 @@ install_binary() { echo '# added by Dria Compute Launcher' >> "$CONFIG_FILE" echo "export PATH=\"${DRIA_INSTALL_DIR}:\$PATH\"" >> "$CONFIG_FILE" else - print_step "Config file $CONFIG_FILE not found" - print_step "Manually add the directory to your shell config:" - print_step "export PATH=\"${DRIA_INSTALL_DIR}:\$PATH\"" + print_step "Creating config file for your shell: $CONFIG_FILE" + touch "$CONFIG_FILE" + echo '# added by Dria Compute Launcher' >> "$CONFIG_FILE" + echo "export PATH=\"${DRIA_INSTALL_DIR}:\$PATH\"" >> "$CONFIG_FILE" fi } diff --git a/src/commands/check.rs b/src/commands/check.rs new file mode 100644 index 0000000..98d4d86 --- /dev/null +++ b/src/commands/check.rs @@ -0,0 +1,35 @@ +use eyre::Context; +use serde::Deserialize; +use std::collections::HashMap; +use prettytable::{Table, row, cell}; + +#[derive(Deserialize)] +struct ApiResponse { + models: HashMap, +} + +pub async fn check_command() -> eyre::Result<()> { + let url = "https://dkn.dria.co/dashboard/supply/v0/models/executions/per-model/last-week"; + + let response = reqwest::get(url) + .await + .wrap_err("Could not make request")?; + + let data: ApiResponse = response + .json() + .await + .wrap_err("Could not parse body")?; + + let mut models: Vec<(&String, &usize)> = data.models.iter().collect(); + models.sort_by(|a, b| b.1.cmp(a.1)); + let mut table = Table::new(); + table.add_row(row!["Model Name", "Tasks"]); + + for (model, task) in models { + table.add_row(row![model, task]); + } + + table.printstd(); + + Ok(()) +} diff --git a/src/commands/mod.rs b/src/commands/mod.rs index 054247a..6bf1161 100644 --- a/src/commands/mod.rs +++ b/src/commands/mod.rs @@ -28,6 +28,10 @@ pub use uninstall::uninstall_launcher; mod points; pub use points::show_points; +mod check; +pub use check::check_command; + + /// Launcher commands. #[derive(clap::Subcommand)] pub enum Commands { @@ -41,6 +45,8 @@ pub enum Commands { Referrals, /// Show your $DRIA points. Points, + /// Fetch and display model usage statistics in the last 7 days + Check, /// Uninstall the launcher & its files. Uninstall, /// Show information about the current environment. diff --git a/src/main.rs b/src/main.rs index da00d6f..468e862 100644 --- a/src/main.rs +++ b/src/main.rs @@ -111,6 +111,7 @@ async fn main() -> eyre::Result<()> { .await; } Commands::Referrals => commands::handle_referrals().await?, + Commands::Check => commands::check_command().await?, }; Ok(())