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