Added ability to set custom memory location (#40)

+1 -1

scripts/install.sh

··· 1 1 #!/usr/bin/env bash 2 2 set -euo pipefail 3 3 4 - ENV="prod" # prod is another env, try taking it from github env 4 + ENV="dev" # prod is another env, try taking it from github env 5 5 REPO="tilesprivacy/tiles" 6 6 # VERSION="${TILES_VERSION:-latest}" 7 7 VERSION="0.3.0"

+13

tiles/src/commands/mod.rs

··· 1 1 // Module that handles CLI commands 2 2 3 + use owo_colors::OwoColorize; 3 4 use tiles::runtime::Runtime; 5 + use tiles::utils::config::set_memory_path; 4 6 use tiles::{core::health, runtime::RunArgs}; 5 7 6 8 pub async fn run(runtime: &Runtime, run_args: RunArgs) { 7 9 runtime.run(run_args).await; 8 10 } 9 11 12 + pub fn set_memory(path: &str) { 13 + match set_memory_path(path) { 14 + Ok(msg) => { 15 + println!("{}", msg.green()); 16 + } 17 + Err(err) => { 18 + let error_msg = format!("Error setting memory path due to {:?}", err); 19 + println!("{}", error_msg.red()); 20 + } 21 + } 22 + } 10 23 pub fn check_health() { 11 24 health::check_health(); 12 25 }

+20

tiles/src/main.rs

··· 22 22 flags: RunFlags, 23 23 }, 24 24 25 + /// Configure your memory 26 + Memory(MemoryArgs), 27 + 25 28 /// Checks the status of dependencies 26 29 Health, 27 30 ··· 55 58 /// Stops the daemon py server 56 59 Stop, 57 60 } 61 + 62 + #[derive(Debug, Args)] 63 + #[command(args_conflicts_with_subcommands = true)] 64 + #[command(flatten_help = true)] 65 + struct MemoryArgs { 66 + #[command(subcommand)] 67 + command: MemoryCommands, 68 + } 69 + #[derive(Debug, Subcommand)] 70 + enum MemoryCommands { 71 + /// Set Path for the memory 72 + SetPath { path: String }, 73 + } 74 + 58 75 #[tokio::main(flavor = "current_thread")] 59 76 pub async fn main() -> Result<(), Box<dyn Error>> { 60 77 let cli = Cli::parse(); ··· 77 94 Some(ServerCommands::Start) => commands::start_server(&runtime).await, 78 95 Some(ServerCommands::Stop) => commands::stop_server(&runtime).await, 79 96 _ => println!("Expected start or stop"), 97 + }, 98 + Commands::Memory(memory) => match memory.command { 99 + MemoryCommands::SetPath { path } => commands::set_memory(path.as_str()), 80 100 }, 81 101 } 82 102 Ok(())

+94 -112

tiles/src/runtime/mlx.rs

··· 1 1 use crate::runtime::RunArgs; 2 + use crate::utils::config::{ 3 + create_default_memory_folder, get_config_dir, get_default_memory_path, get_memory_path, 4 + get_server_dir, set_memory_path, 5 + }; 2 6 use crate::utils::hf_model_downloader::*; 3 7 use anyhow::{Context, Result}; 4 8 use futures_util::StreamExt; ··· 11 15 use rustyline::validate::Validator; 12 16 use rustyline::{Config, Editor, Helper}; 13 17 use serde_json::{Value, json}; 18 + use std::fs; 14 19 use std::fs::File; 15 - use std::path::PathBuf; 16 - use std::process::Command; 17 20 use std::process::Stdio; 18 21 use std::time::Duration; 19 - use std::{env, fs}; 22 + use std::{io, process::Command}; 20 23 use tilekit::modelfile::Modelfile; 21 24 use tokio::time::sleep; 22 25 pub struct MLXRuntime {} ··· 41 44 42 45 pub async fn run(&self, run_args: super::RunArgs) { 43 46 const DEFAULT_MODELFILE: &str = "FROM driaforall/mem-agent-mlx-4bit"; 44 - 45 - // Parse modelfile 47 + //Parse modelfile 46 48 let modelfile_parse_result = if let Some(modelfile_str) = &run_args.modelfile_path { 47 49 tilekit::modelfile::parse_from_file(modelfile_str.as_str()) 48 50 } else { ··· 59 61 60 62 let model = modelfile.from.as_ref().unwrap(); 61 63 if model.starts_with("driaforall/mem-agent") { 62 - let _res = run_model_with_server(self, modelfile, &run_args).await; 64 + let _res = run_model_with_server(self, modelfile, &run_args) 65 + .await 66 + .inspect_err(|e| eprintln!("Failed to run the model due to {e}")); 63 67 } else { 64 68 run_model_by_sub_process(modelfile); 65 69 } ··· 258 262 mlx_runtime: &MLXRuntime, 259 263 modelfile: Modelfile, 260 264 run_args: &RunArgs, 261 - ) -> reqwest::Result<()> { 265 + ) -> Result<()> { 262 266 if !cfg!(debug_assertions) { 263 - let _res = mlx_runtime.start_server_daemon().await; 267 + let _ = mlx_runtime.start_server_daemon().await.inspect_err(|e| { 268 + eprintln!("Failed to start daemon server due to {:?}", e); 269 + }); 264 270 let _ = wait_until_server_is_up().await; 265 271 } 266 272 // loading the model from mem-agent via daemon server 267 - let memory_path = get_memory_path() 268 - .context("Retrieving memory_path failed") 269 - .unwrap(); 273 + let memory_path = get_or_set_memory_path().context("Setting/Retrieving memory_path failed")?; 270 274 let modelname = modelfile.from.as_ref().unwrap(); 271 275 match load_model(modelname, &memory_path).await { 272 276 Ok(_) => start_repl(mlx_runtime, modelname, run_args).await, 273 - Err(err) => println!("{}", err), 277 + Err(err) => return Err(anyhow::anyhow!(err)), 274 278 } 275 279 Ok(()) 276 280 } 277 281 282 + fn get_or_set_memory_path() -> Result<String> { 283 + match get_memory_path() { 284 + Ok(memory_path) => Ok(memory_path), 285 + Err(_err) => { 286 + let stdin = io::stdin(); 287 + let default_memory_pathbuf = get_default_memory_path()?; 288 + let mut default_memory = default_memory_pathbuf 289 + .to_str() 290 + .ok_or_else(|| anyhow::anyhow!("Invalid path"))?; 291 + let mut chose_yes = false; 292 + 293 + println!( 294 + "{}", 295 + format!( 296 + "Default Memory location will be set at {:?}\n", 297 + default_memory 298 + ) 299 + .yellow() 300 + ); 301 + println!("You can always change the location with `tiles memory set-path <PATH>`\n"); 302 + println!("Do you want to add a custom memory location right now instead? [Y/N]"); 303 + let mut input = String::new(); 304 + loop { 305 + input.clear(); 306 + stdin.read_line(&mut input)?; 307 + input = input.trim().to_owned(); 308 + if (input == "Y" || input == "y") || chose_yes { 309 + if !chose_yes { 310 + chose_yes = true; 311 + println!("Add the path for your custom memory"); 312 + continue; 313 + } 314 + match set_memory_path(input.as_str()) { 315 + Ok(msg) => { 316 + default_memory = input.as_str(); 317 + println!("{}", msg.green()); 318 + println!( 319 + "You can always change the location with `tiles memory set-path <PATH>`\n" 320 + ); 321 + break; 322 + } 323 + Err(err) => { 324 + let error_msg = 325 + format!("Try again, Error setting memory path due to {:?}", err); 326 + println!("{}", error_msg.red()); 327 + continue; 328 + } 329 + } 330 + } else { 331 + create_default_memory_folder()?; 332 + match set_memory_path(default_memory) { 333 + Ok(msg) => { 334 + println!("{}", msg.green()); 335 + println!( 336 + "You can always change the location with `tiles memory set-path <PATH>`\n" 337 + ); 338 + break; 339 + } 340 + Err(err) => { 341 + let error_msg = format!("Error setting memory path due to {:?}", err); 342 + println!("{}", error_msg.red()); 343 + return Err(anyhow::anyhow!("Error setting default memory path")); 344 + } 345 + } 346 + } 347 + } 348 + Ok(default_memory.to_owned()) 349 + } 350 + } 351 + } 352 + 278 353 async fn start_repl(mlx_runtime: &MLXRuntime, modelname: &str, run_args: &RunArgs) { 279 354 println!("Running in interactive mode"); 280 355 281 - // Setup rustyline editor with hint support 282 356 let config = Config::builder().auto_add_history(true).build(); 283 357 let mut editor = Editor::<TilesHinter, DefaultHistory>::with_config(config).unwrap(); 284 358 editor.set_helper(Some(TilesHinter)); 285 359 286 - // TODO: Handle "enter" key press or any key press when repl is processing an input 287 360 loop { 288 361 let readline = editor.readline(">>> "); 289 362 let input = match readline { ··· 298 371 } 299 372 }; 300 373 301 - // Handle slash commands 302 374 match handle_slash_command(&input, modelname) { 303 375 SlashCommand::Continue => continue, 304 376 SlashCommand::Exit => { ··· 311 383 SlashCommand::NotACommand => {} 312 384 } 313 385 314 - // Skip empty input 315 386 if input.is_empty() { 316 387 continue; 317 388 } 318 - 319 - // Send to model 320 389 let mut remaining_count = run_args.relay_count; 321 390 let mut g_reply: String = "".to_owned(); 322 391 let mut python_code: String = "".to_owned(); ··· 358 427 } 359 428 } 360 429 361 - async fn load_model(model_name: &str, memory_path: &str) -> Result<(), String> { 430 + async fn load_model(model_name: &str, memory_path: &str) -> Result<()> { 362 431 let client = Client::new(); 363 432 let body = json!({ 364 433 "model": model_name, 365 434 "memory_path": memory_path 366 435 }); 367 436 368 - //TODO: Fix the unwrap here 369 437 let res = client 370 438 .post("http://127.0.0.1:6969/start") 371 439 .json(&body) 372 440 .send() 373 - .await 374 - .unwrap(); 441 + .await?; 375 442 match res.status() { 376 443 StatusCode::OK => Ok(()), 377 444 StatusCode::NOT_FOUND => { ··· 381 448 println!("\nDownloading completed \n"); 382 449 Ok(()) 383 450 } 384 - Err(err) => Err(err), 451 + Err(err) => Err(anyhow::anyhow!(format!("Download failed due to {:?}", err))), 385 452 } 386 453 } 387 - _ => { 388 - println!("err {:?}", res); 389 - Err(format!( 390 - "Failed to load model {} due to {:?}", 391 - model_name, res 392 - )) 393 - } 454 + _ => Err(anyhow::anyhow!(format!( 455 + "Failed to load model {} due to {:?}", 456 + model_name, res 457 + ))), 394 458 } 395 459 } 396 460 ··· 446 510 } 447 511 448 512 fn convert_to_chat_response(content: &str) -> ChatResponse { 449 - // content.split() 450 513 ChatResponse { 451 514 reply: extract_reply(content), 452 515 code: extract_python(content), ··· 470 533 list_b[0].to_owned() 471 534 } else { 472 535 "".to_owned() 473 - } 474 - } 475 - 476 - // fn extract_think(content: &str) -> String { 477 - // if content.contains("<think>") && content.contains("</think>") { 478 - // let list_a = content.split("<think>").collect::<Vec<&str>>(); 479 - // let list_b = list_a[1].split("</think>").collect::<Vec<&str>>(); 480 - // list_b[0].to_owned() 481 - // } else if content.contains("</think") { 482 - // let list_a = content.split("</think>").collect::<Vec<&str>>(); 483 - // list_a[0].to_owned() 484 - // } else { 485 - // "".to_owned() 486 - // } 487 - // } 488 - 489 - fn get_memory_path() -> Result<String> { 490 - let tiles_config_dir = get_config_dir()?; 491 - let tiles_data_dir = get_data_dir()?; 492 - let mut is_memory_path_found: bool = false; 493 - let mut memory_path: String = String::from(""); 494 - if tiles_config_dir.is_dir() 495 - && let Ok(content) = fs::read_to_string(tiles_config_dir.join(".memory_path")) 496 - { 497 - memory_path = content; 498 - is_memory_path_found = true; 499 - } 500 - 501 - if is_memory_path_found { 502 - Ok(memory_path) 503 - } else { 504 - let memory_path = tiles_data_dir.join("memory"); 505 - fs::create_dir_all(&memory_path).context("Failed to create tiles memory directory")?; 506 - fs::create_dir_all(&tiles_config_dir).context("Failed to create tiles config directory")?; 507 - fs::write( 508 - tiles_config_dir.join(".memory_path"), 509 - memory_path.to_str().unwrap(), 510 - ) 511 - .context("Failed to write the default path to .memory_path")?; 512 - Ok(memory_path.to_string_lossy().to_string()) 513 - } 514 - } 515 - 516 - fn get_server_dir() -> Result<PathBuf> { 517 - if cfg!(debug_assertions) { 518 - let base_dir = env::current_dir().context("Failed to fetch CURRENT_DIR")?; 519 - Ok(base_dir.join("server")) 520 - } else { 521 - let home_dir = env::home_dir().context("Failed to fetch $HOME")?; 522 - let data_dir = match env::var("XDG_DATA_HOME") { 523 - Ok(val) => PathBuf::from(val), 524 - Err(_err) => home_dir.join(".local/share"), 525 - }; 526 - Ok(data_dir.join("tiles/server")) 527 - } 528 - } 529 - fn get_config_dir() -> Result<PathBuf> { 530 - if cfg!(debug_assertions) { 531 - let base_dir = env::current_dir().context("Failed to fetch CURRENT_DIR")?; 532 - Ok(base_dir.join(".tiles_dev/tiles")) 533 - } else { 534 - let home_dir = env::home_dir().context("Failed to fetch $HOME")?; 535 - let config_dir = match env::var("XDG_CONFIG_HOME") { 536 - Ok(val) => PathBuf::from(val), 537 - Err(_err) => home_dir.join(".config"), 538 - }; 539 - Ok(config_dir.join("tiles")) 540 - } 541 - } 542 - 543 - fn get_data_dir() -> Result<PathBuf> { 544 - if cfg!(debug_assertions) { 545 - let base_dir = env::current_dir().context("Failed to fetch CURRENT_DIR")?; 546 - Ok(base_dir.join(".tiles_dev/tiles")) 547 - } else { 548 - let home_dir = env::home_dir().context("Failed to fetch $HOME")?; 549 - let data_dir = match env::var("XDG_DATA_HOME") { 550 - Ok(val) => PathBuf::from(val), 551 - Err(_err) => home_dir.join(".local/share"), 552 - }; 553 - Ok(data_dir.join("tiles")) 554 536 } 555 537 } 556 538

+102

tiles/src/utils/config.rs

··· 1 + // Configuration related stuff 2 + 3 + use anyhow::{Context, Result}; 4 + use std::path::PathBuf; 5 + use std::str::FromStr; 6 + use std::{env, fs}; 7 + 8 + pub fn set_memory_path(path: &str) -> Result<String> { 9 + let path_buf = PathBuf::from_str(path)?; 10 + if path_buf.try_exists()? { 11 + let tiles_config_dir = get_config_dir()?; 12 + let tiles_config_memory_path = tiles_config_dir.join(".memory_path"); 13 + if tiles_config_memory_path.try_exists()? { 14 + fs::write(tiles_config_memory_path, path)?; 15 + } else { 16 + fs::create_dir_all(&tiles_config_dir) 17 + .context("Failed to create tiles config directory")?; 18 + fs::write(tiles_config_memory_path, path) 19 + .context("Failed to write the memory path to .memory_path")?; 20 + } 21 + } else { 22 + return Err(anyhow::anyhow!(format!( 23 + "Not a valid path {}", 24 + path_buf.to_str().unwrap() 25 + ))); 26 + } 27 + Ok(format!( 28 + "Memory path set successfully at {}", 29 + path_buf.to_str().unwrap() 30 + )) 31 + } 32 + 33 + pub fn get_memory_path() -> Result<String> { 34 + let tiles_config_dir = get_config_dir()?; 35 + let mut is_memory_path_found: bool = false; 36 + let mut memory_path: String = String::from(""); 37 + if tiles_config_dir.is_dir() 38 + && let Ok(content) = fs::read_to_string(tiles_config_dir.join(".memory_path")) 39 + { 40 + memory_path = content; 41 + is_memory_path_found = true; 42 + } 43 + 44 + if is_memory_path_found { 45 + Ok(memory_path) 46 + } else { 47 + Err(anyhow::anyhow!(format!("NOT SET"))) 48 + } 49 + } 50 + 51 + pub fn get_default_memory_path() -> Result<PathBuf> { 52 + let tiles_data_dir = get_data_dir()?; 53 + let memory_path = tiles_data_dir.join("memory"); 54 + Ok(memory_path) 55 + } 56 + 57 + pub fn create_default_memory_folder() -> Result<PathBuf> { 58 + let memory_path = get_default_memory_path()?; 59 + fs::create_dir_all(&memory_path).context("Failed to create tiles memory directory")?; 60 + Ok(memory_path) 61 + } 62 + 63 + pub fn get_server_dir() -> Result<PathBuf> { 64 + if cfg!(debug_assertions) { 65 + let base_dir = env::current_dir().context("Failed to fetch CURRENT_DIR")?; 66 + Ok(base_dir.join("server")) 67 + } else { 68 + let home_dir = env::home_dir().context("Failed to fetch $HOME")?; 69 + let data_dir = match env::var("XDG_DATA_HOME") { 70 + Ok(val) => PathBuf::from(val), 71 + Err(_err) => home_dir.join(".local/share"), 72 + }; 73 + Ok(data_dir.join("tiles/server")) 74 + } 75 + } 76 + pub fn get_config_dir() -> Result<PathBuf> { 77 + if cfg!(debug_assertions) { 78 + let base_dir = env::current_dir().context("Failed to fetch CURRENT_DIR")?; 79 + Ok(base_dir.join(".tiles_dev/tiles")) 80 + } else { 81 + let home_dir = env::home_dir().context("Failed to fetch $HOME")?; 82 + let config_dir = match env::var("XDG_CONFIG_HOME") { 83 + Ok(val) => PathBuf::from(val), 84 + Err(_err) => home_dir.join(".config"), 85 + }; 86 + Ok(config_dir.join("tiles")) 87 + } 88 + } 89 + 90 + pub fn get_data_dir() -> Result<PathBuf> { 91 + if cfg!(debug_assertions) { 92 + let base_dir = env::current_dir().context("Failed to fetch CURRENT_DIR")?; 93 + Ok(base_dir.join(".tiles_dev/tiles")) 94 + } else { 95 + let home_dir = env::home_dir().context("Failed to fetch $HOME")?; 96 + let data_dir = match env::var("XDG_DATA_HOME") { 97 + Ok(val) => PathBuf::from(val), 98 + Err(_err) => home_dir.join(".local/share"), 99 + }; 100 + Ok(data_dir.join("tiles")) 101 + } 102 + }

+1

tiles/src/utils/mod.rs

··· 1 + pub mod config; 1 2 pub mod hf_model_downloader;

Configure Feed

Configure Feed