use serde_json::Value; use std::fs::File; use std::io::prelude::*; use std::io::stdin; use std::sync::mpsc::channel; use std::sync::{ Arc, Barrier }; use std::path::Path; mod data; mod node; /** * Parse the JSON file and store the result in the parsed variable * * @param arg: String file path * @param parsed: &mut Value variable to store the result * @param debug: bool debug flag * * @return std::io::Result<()> */ fn parse_json(arg: String, parsed: &mut Value, debug: bool) -> std::io::Result<()> { if arg == "exit" { return Err(std::io::Error::new(std::io::ErrorKind::Other, "exit")); } let mut _file: Option<File> = None; match File::open(arg) { Ok(opened_file) => _file = Some(opened_file), Err(_) => { let _ = parse_json(get_input(Some("Invalid path")), parsed, debug); return Ok(()); } } if let Some(ref mut actual_file) = _file { let mut content = String::new(); actual_file.read_to_string(&mut content)?; *parsed = serde_json::from_str(&content)?; } if debug == true { println!( "|----------------parse_json--------------|\n{}", parsed.to_string() ); } Ok(()) } /** * Print the help message */ fn print_help() { println!("|----------------HELP-------------------|"); println!("<node> <command> => send command to node"); println!("exit => exit program"); println!("list => list avaible nodes"); println!("<node> quit => quit node"); println!("command:\n send <message> => send message to llm\n quit => quit node"); println!("help => display this message"); println!("----------------------------------------"); } /** * Get the input from the user * * @return String */ fn get_input(msg: Option<&str>) -> String { loop { println!("{}", msg.unwrap_or("").to_string()); let mut input = String::new(); match stdin().read_line(&mut input) { Ok(_) => { return input.trim().to_string(); } Err(_) => { println!("Error reading input"); return String::from(""); } } } } /** * Main loop * * @param barrier: Arc<Barrier> barrier to sync the threads * @param nodes: &mut Vec<Node> vector of nodes * @param debug: bool debug flag * * @return std::io::Result<()> */ fn main_loop(barrier: Arc<Barrier>, nodes: &mut Vec<data::Node>, _flags: data::Flags) -> std::io::Result<()> { // init input barrier.wait(); let mut input = String::new(); print_help(); println!("nodes :"); for node in nodes.iter() { println!(" {}", node.id); } println!("----------------------------------------"); let mut did_something = false; loop { match stdin().read_line(&mut input) { Ok(_) => { match input.trim() { "exit" => { println!("exiting the program"); for node in nodes.iter() { node.chan.send(Value::String("quit".to_string())).unwrap(); } return Ok(()); } "list" => { for node in nodes.iter() { println!("{}", node.id); } } _ => { // check if the input is a node command let mut index_to_remove = None; for (index, node) in nodes.iter().enumerate() { let (target, mut command) = input.trim().split_at(input.trim().find(" ").unwrap_or(0)); command = command.trim(); if target == node.id { did_something = true; node.chan.send(Value::String(command.to_string())).unwrap(); // check if the command is quit, need to remove the node if true if command == "quit" { index_to_remove = Some(index); } } } if let Some(index) = index_to_remove { nodes.remove(index); } if did_something == false { print_help(); } else { did_something = false; } } } input.clear(); } Err(_) => print_help(), } } } /** * Get the arguments from the command line * * @param args: std::env::Args arguments * @param arg: &mut String file path * @param flags: &mut flags flags * * @return std::io::Result<()> */ fn get_args(args: std::env::Args, arg: &mut String, flags: &mut data::Flags) -> std::io::Result<()> { let mut got_path = false; let mut last = "none"; for args_it in args.skip(1) { if last == "-m" { flags.model_name = args_it.to_string(); last = "none"; continue; } match args_it.as_str() { "-d" => { println!("Debug mode enabled"); flags.debug = true; }, "-m" => { flags.model = true; last = "-m"; }, _ => { if Path::new(&args_it).exists() { *arg = args_it.to_string(); got_path = true; } else { println!("Invalid argument: {}", args_it); return Err(std::io::Error::new(std::io::ErrorKind::Other, "Invalid argument")); } } } } let mut _path: String = String::new(); while got_path == false { _path = get_input(Some("Enter the path to the JSON file")); if _path == "exit" { return Err(std::io::Error::new(std::io::ErrorKind::Other, "exit")); } else if Path::new(&_path).exists() { got_path = true; } else { println!("Invalid path"); } } Ok(()) } fn main() -> std::io::Result<()> { let mut arg: String = String::new(); let mut flags: data::Flags = data::Flags { debug: false, model: false, model_name: String::new(), }; match get_args(std::env::args(), &mut arg, &mut flags){ Ok(_) => (), Err(e) => { println!("Error getting arguments: {}", e.to_string()); return Ok(()); } } if flags.debug == true { println!("|----------------get_arg-------------------|"); println!("arg: {}", arg); println!("model: {}", flags.model_name); } let mut parsed = serde_json::Value::Null; let mut nodes: Vec<data::Node> = Vec::new(); match parse_json(arg, &mut parsed, flags.debug) { Ok(_) => (), Err(e) => { println!("Error parsing JSON file: {}", e.to_string()); return Ok(()); } } let barrier = Arc::new(Barrier::new( parsed.as_object().expect("Expected a JSON object").len() + 1, )); // +1 for the main thread if let Value::Object(map) = &parsed { for (key, value) in map { let (tx, rx) = channel(); let new_node: data::Node = data::Node { id: key.to_string(), chan: tx.clone(), }; if let Value::Array(_array) = value { node::make_node(Arc::clone(&barrier), key.to_string() , value.clone(), rx, flags.clone()); } nodes.push(new_node); } } main_loop(Arc::clone(&barrier), &mut nodes, flags)?; barrier.wait(); Ok(()) } //https://github.com/pepperoni21/ollama-rs