node.rs

use ollama_rs::generation::completion::request::GenerationRequest;
use std::process::Command;
use serde_json::Value;
use std::net::{
    SocketAddr, 
    UdpSocket
};
use std::{
    thread, 
    time
};
use ollama_rs::*;
use tokio::runtime::Runtime;
use chrono::Utc;
use std::sync::{
    Barrier, 
    Arc
};

use crate::data;

const BUFF_SIZE: usize = 4096;
const MODEL: &str = "mistral:latest";
const LOCALHOST: &str = "http://localhost";

/**
 * Make a node
 * 
 * @param barrier: Arc<Barrier> barrier to sync the threads
 * @param node_id: String node id
 * @param array: Value array of the node
 * @param rx: std::sync::mpsc::Receiver<Value> receiver to comunicate from the master thread
 * @param flags: data::Flags flags
 * 
 * @return 
 */
pub fn make_node(barrier: Arc<Barrier>, node_id: String , array: Value, rx: std::sync::mpsc::Receiver<Value>, flags: data::Flags) {

    if flags.debug == true {
        println!("|----------------make_node: {}----------------|", node_id);
    }

    let _child = thread::spawn(move || {
        let llm_string = array[1].get("llm").unwrap().as_str().unwrap();

        if flags.debug == true {
            println!("{} : llm -> {}, {}", node_id, llm_string, llm_string.split(":").collect::<Vec<&str>>()[1]);
        }
        
        let _ =  if cfg!(target_os = "windows") {
            Command::new("cmd")
                .args(&["/C", &format!("docker run -d -v ollama:/root/.ollama -p {}:11434 --name ollama_{} ollama/ollama", llm_string.split(":").collect::<Vec<&str>>()[1], node_id)])
                .output()
                .expect("failed to execute process")
        } else {
            Command::new("sh")
                    .arg("-c")
                    .arg(&format!("docker run -d -v ollama:/root/.ollama -p {}:11434 --name ollama_{} ollama/ollama", llm_string.split(":").collect::<Vec<&str>>()[1], node_id))
                    .output()
                    .expect("failed to execute process")
        };

        println!("{} -> arbitrary wait for the container to start", node_id);
        thread::sleep(time::Duration::from_secs(5));

        let mut _model = String::new();
        if flags.model == true {
            _model = flags.model_name.clone();
        } else {
            _model = MODEL.to_string();
        }

        let _ =  if cfg!(target_os = "windows") {
            Command::new("cmd")
                .args(&["/C", &format!("docker exec --it ollama_{} ollama run {}", node_id, _model)])
                .output()
                .expect("failed to execute process")
        } else {
            Command::new("sh")
                    .arg("-c")
                    .arg(&format!("docker exec --it ollama_{} ollama run {}", node_id, _model))
                    .output()
                    .expect("failed to execute process")
        };

        println!("{} -> arbitrary wait for the model to start", node_id);
        thread::sleep(time::Duration::from_secs(5));

        let node_socket = UdpSocket::bind(array[0].get("node").unwrap().as_str().unwrap()).expect("Could not bind node address");
        let mut comm_vec: Vec<SocketAddr> = Vec::new();
        
        let llm = Ollama::new(LOCALHOST.to_string(), llm_string.split(":").collect::<Vec<&str>>()[1].parse::<u16>().unwrap());

        node_socket.set_nonblocking(true).expect("Could not set non-blocking mode to node socket");

        if let Some(Value::Array(comm_arr)) = array[2].get("neighbours") {
            for comm in comm_arr {
                if let Some(comm_str) = comm.as_str() {
                    if let Ok(comm_addr) = comm_str.parse::<std::net::SocketAddr>() {
                        comm_vec.push(comm_addr);
                    } else {
                        println!("{} : Invalid comm address: {}", node_id, comm_str);
                    }
                } else {
                    println!("{} : Invalid comm address: {:?}", node_id, comm);
                }
            }
        }

        if flags.debug == true {
            println!("{} : node -> {}", node_id, node_socket.local_addr().unwrap());
            for comm in array[2].get("neighbours").iter() {
                println!("{} : comm -> {}", node_id, comm);
            }
        }

        barrier.wait();

        // node_loop(barrier, node_id, rx, node_socket, llm, array[2].get("neighbours"), debug).unwrap();
        node_loop(barrier, node_id, rx, node_socket, llm, comm_vec, flags.debug).unwrap();

    });
}

/*
* Node loop
*
* @param barrier: Arc<Barrier> barrier to sync the threads
* @param node_id: String node id
* @param rx: std::sync::mpsc::Receiver<Value> receiver to comunicate from the master thread
* @param node_socket: std::net::UdpSocket node socket
* @param llm: Ollama Ollama instance
* @param comm_vec: Vec<SocketAddr> vector of neighbours
* @param debug: bool debug flag
*
* @return std::io::Result<()>
*/
fn node_loop( barrier: Arc<Barrier>, node_id: String, rx: std::sync::mpsc::Receiver<Value>, node_socket: std::net::UdpSocket, llm: Ollama,comm_vec: Vec<SocketAddr>, debug: bool) -> std::io::Result<()> {
    let mut msg_received = false;
    let mut buf = [0; BUFF_SIZE];
    let mut data: data::Message = data::Message { id: "init".to_string(), message: "".to_string() };
    let mut response: data::Message = data::Message { id: "".to_string(), message: "".to_string() };
    let start_time = Utc::now().timestamp();
    let rt = Runtime::new().unwrap();

    loop {
        // listen to master
        let elapsed = Utc::now().timestamp() - start_time;
        match rx.try_recv() {
            Ok(Value::String(msg)) => {
                if debug == true {
                    println!("{} -> {}", node_id, msg);
                }

                let (first, rest) = match msg.split_whitespace().next() {
                    Some(first) => (first, msg.trim_start_matches(first).trim_start()),
                    None => continue,
                };

                if debug == true {
                    println!("{} (first, rest) -> {} {}", node_id, first, rest);
                }

                match first.trim().replace(" ", "").as_str() {
                    "send" => {
                        msg_received = true;
                        data.message = rest.to_string();
                        
                        if debug == true {
                            println!("{} : send received -> {}", node_id, data.message);
                        }

                    },
                    "quit" => {
                        drop(node_socket);
                        if cfg!(target_os = "windows") {
                            match Command::new("cmd")
                                .args(&["/C", &format!("docker rm -f /ollama_{}", node_id)])
                                .output() {
                                    Ok(_) => {},
                                    Err(e) => println!("{} : Error: {}", node_id, e)
                                }
                        } else {
                            match Command::new("sh")
                                .arg("-c")
                                .arg(&format!("docker rm -f /ollama_{}", node_id))
                                .output() {
                                    Ok(_) => {},
                                    Err(e) => println!("{} : Error: {}", node_id, e)
                                }
                        };

                        println!("child exited");
                        
                        barrier.wait();

                        return Ok(());
                    },
                    _ => {
                        println!("{} : Invalid command: {}", node_id, first);
                    }
                    
                }
            },
            Ok(Value::Null) => {
                continue;
            },
            Ok(Value::Bool(_)) => {
                println!("Bool value not supported yet!");
            },
            Ok(Value::Number(_)) => {
                println!("Number value not supported yet!");
            },
            Ok(Value::Array(_)) => {
                println!("Array value not supported yet!");
            },
            Ok(Value::Object(_)) => {
                println!("Object value not supported yet!");
            },
            Err(std::sync::mpsc::TryRecvError::Empty) => {
            },
            Err(e) => println!("Error node_loop: {}", e)
        }

        match node_socket.recv_from(&mut buf) {
            Ok(_) => {
                let json_end = buf.iter().position(|&b| b == b'}');
                if let Some(end) = json_end {
                    let json_slice = &buf[..=end];
                    data = serde_json::from_slice(json_slice).unwrap();
                    if debug == true {
                        println!("{:?}:{} received -> {}", elapsed, node_id, data.message);
                    }
                    msg_received = true;
                } else {
                    println!("Invalid JSON data");
                }
            },
            Err(e) => {
                if e.kind() != std::io::ErrorKind::WouldBlock {
                    println!("{} : Error: {}", node_id, e);
                }
            }
        }
        
        rt.block_on(async {
            //listen to neighbours
            if msg_received == true {
                if debug == true {
                    println!("{:?}:{} request -> {}", elapsed, node_id, data.message);
                }
                response.message = llm.generate(GenerationRequest::new(MODEL.to_string(), data.message.clone())).await.clone().unwrap().response;
                println!("{:?}: {} response -> {}", elapsed, node_id, response.message);
                response.id = node_id.clone();
                if debug == true {
                    println!("{:?}:{} response -> {}", elapsed, node_id, response.message);
                }
                
                let send = serde_json::to_string(&response).unwrap();
                for comm in comm_vec.iter() {
                    let comm_str = comm.to_string();
                    let addr = comm_str.trim_matches(|c| c == '[' || c == ']' || c == '\"');
                    println!("{:?}:{} sending to -> {}", elapsed, node_id, addr);
                    match node_socket.send_to(send.as_bytes(), addr){
                        Ok(_) => {
                            if debug == true {
                                println!("{:?}:{} sent -> {}", elapsed, node_id, send);
                            }
                        },
                        Err(e) => {
                            println!("{} : Error: {}", node_id, e);
                        }
                    }
                }
            }
            msg_received = false;
        });
    }
}