added the count flag and made the code scalable to add future flags if needed (#35)

Author: kshitijgetsac

tiles/src/commands/mod.rs

@@ -1,26 +1,10 @@
 // Module that handles CLI commands
 
-use anyhow::Result;
-use tilekit::{modelfile, modelfile::Modelfile};
 use tiles::runtime::Runtime;
 use tiles::{core::health, runtime::RunArgs};
-const DEFAULT_MODELFILE: &str = "
-  FROM driaforall/mem-agent-mlx-4bit 
-";
 
-pub async fn run(runtime: &Runtime, modelfile: Option<String>) {
-    let modelfile_parse_result: Result<Modelfile, String> = if let Some(modelfile_str) = modelfile {
-        modelfile::parse_from_file(modelfile_str.as_str())
-    } else {
-        modelfile::parse(DEFAULT_MODELFILE)
-    };
-    match modelfile_parse_result {
-        Ok(modelfile) => {
-            let run_args = RunArgs { modelfile };
-            runtime.run(run_args).await;
-        }
-        Err(_err) => println!("Invalid Modelfile"),
-    }
+pub async fn run(runtime: &Runtime, run_args: RunArgs) {
+    runtime.run(run_args).await;
 }
 
 pub fn check_health() {

tiles/src/main.rs

@@ -1,7 +1,7 @@
 use std::error::Error;
 
 use clap::{Args, Parser, Subcommand};
-use tiles::runtime::build_runtime;
+use tiles::runtime::{build_runtime, RunArgs};
 mod commands;
 #[derive(Debug, Parser)]
 #[command(name = "tiles")]
@@ -14,7 +14,13 @@ struct Cli {
 #[derive(Subcommand, Debug)]
 enum Commands {
     /// Runs the given Modelfile (runs the default model if none passed)
-    Run { modelfile_path: Option<String> },
+    Run {
+        /// Path to the Modelfile (uses default model if not provided)
+        modelfile_path: Option<String>,
+
+        #[command(flatten)]
+        flags: RunFlags,
+    },
 
     /// Checks the status of dependencies
     Health,
@@ -23,6 +29,17 @@ enum Commands {
     Server(ServerArgs),
 }
 
+#[derive(Debug, Args)]
+struct RunFlags {
+    /// Number of chat retries before giving up (default: 6)
+    #[arg(short = 'r', long, default_value_t = 6)]
+    retry_count: u32,
+
+    // Future flags go here:
+    // #[arg(long, default_value_t = 6969)]
+    // port: u16,
+}
+
 #[derive(Debug, Args)]
 #[command(args_conflicts_with_subcommands = true)]
 #[command(flatten_help = true)]
@@ -44,8 +61,12 @@ pub async fn main() -> Result<(), Box<dyn Error>> {
     let cli = Cli::parse();
     let runtime = build_runtime();
     match cli.command {
-        Commands::Run { modelfile_path } => {
-            commands::run(&runtime, modelfile_path).await;
+        Commands::Run { modelfile_path, flags } => {
+            let run_args = RunArgs {
+                modelfile_path,
+                retry_count: flags.retry_count,
+            };
+            commands::run(&runtime, run_args).await;
         }
         Commands::Health => {
             commands::check_health();

tiles/src/runtime/mlx.rs

@@ -34,11 +34,28 @@ impl MLXRuntime {
     }
 
     pub async fn run(&self, run_args: super::RunArgs) {
-        let model = run_args.modelfile.from.as_ref().unwrap();
+        const DEFAULT_MODELFILE: &str = "FROM driaforall/mem-agent-mlx-4bit";
+
+        // Parse modelfile
+        let modelfile_parse_result = if let Some(modelfile_str) = run_args.modelfile_path {
+            tilekit::modelfile::parse_from_file(modelfile_str.as_str())
+        } else {
+            tilekit::modelfile::parse(DEFAULT_MODELFILE)
+        };
+
+        let modelfile = match modelfile_parse_result {
+            Ok(mf) => mf,
+            Err(_err) => {
+                println!("Invalid Modelfile");
+                return;
+            }
+        };
+
+        let model = modelfile.from.as_ref().unwrap();
         if model.starts_with("driaforall/mem-agent") {
-            let _res = run_model_with_server(self, run_args.modelfile).await;
+            let _res = run_model_with_server(self, modelfile, run_args.retry_count).await;
         } else {
-            run_model_by_sub_process(run_args.modelfile);
+            run_model_by_sub_process(modelfile);
         }
     }
 
@@ -155,6 +172,7 @@ fn run_model_by_sub_process(modelfile: Modelfile) {
 async fn run_model_with_server(
     mlx_runtime: &MLXRuntime,
     modelfile: Modelfile,
+    retry_count: u32,
 ) -> reqwest::Result<()> {
     if !cfg!(debug_assertions) {
         let _res = mlx_runtime.start_server_daemon().await;
@@ -185,12 +203,12 @@ async fn run_model_with_server(
                 break;
             }
             _ => {
-                let mut remaining_count = 6;
+                let mut remaining_count = retry_count;
                 let mut g_reply: String = "".to_owned();
                 let mut python_code: String = "".to_owned();
                 loop {
                     if remaining_count > 0 {
-                        let chat_start = remaining_count == 6;
+                        let chat_start = remaining_count == retry_count;
                         if let Ok(response) = chat(input, modelname, chat_start, &python_code).await
                         {
                             if response.reply.is_empty() {

tiles/src/runtime/mod.rs

@@ -2,12 +2,13 @@
 use crate::runtime::cpu::CPURuntime;
 use crate::runtime::mlx::MLXRuntime;
 use anyhow::Result;
-use tilekit::modelfile::Modelfile;
 pub mod cpu;
 pub mod mlx;
 
 pub struct RunArgs {
-    pub modelfile: Modelfile,
+    pub modelfile_path: Option<String>,
+    pub retry_count: u32,
+    // Future flags go here
 }
 
 pub enum Runtime {