lib/
lib.rs

1use std::fs::{create_dir_all, File, OpenOptions};
2use std::io::{Read, Write};
3use std::path::Path;
4
5use anyhow::{anyhow, ensure, Result};
6use bat::PrettyPrinter;
7use chrono::Local;
8use clap::ArgMatches;
9use colored::Colorize;
10use dialoguer::{theme::ColorfulTheme, Confirm, Input};
11use directories::ProjectDirs;
12use serde_derive::{Deserialize, Serialize};
13use toml::to_string;
14
15struct Project<T: AsRef<str>> {
16    qualifier: T,
17    org: T,
18    app: T,
19}
20
21#[derive(Serialize, Deserialize)]
22pub struct Config {
23    pub host: String,  // Server Addr
24    pub port: u16,     // Server port
25    pub model: String, // Model name
26    pub verbose: bool, // Verbose output following response
27    pub color: bool,   // Color output
28    pub save: bool,    // Autosave conversation
29}
30
31pub enum ContentType {
32    Error,
33    Info,
34    Answer,
35    Exit,
36}
37
38#[derive(Debug)]
39pub enum LogLevel {
40    Debug,
41    Error,
42    Info,
43}
44
45pub enum ProjFiles {
46    Conf,
47    Data,
48    Log,
49}
50
51const PROJECT: Project<&'static str> = Project {
52    qualifier: "io",
53    org: "rtwo",
54    app: "rtwo",
55};
56
57const LOG_FILE: &str = "rtwo.log";
58const CONF_FILE: &str = "rtwo.toml";
59const DB_FILE: &str = "rtwo.db";
60
61pub fn log(lvl: LogLevel, descriptor: &str, msg: &str) -> Result<()> {
62    let log_msg = format!("{:?} {:?} [{}]: {}\n", Local::now(), lvl, descriptor, msg);
63    let log_file = get_project_file(ProjFiles::Log)?;
64    let mut f = match Path::new(&log_file).exists() {
65        true => OpenOptions::new().append(true).open(log_file)?,
66        false => File::create(log_file)?,
67    };
68    f.write_all(log_msg.as_bytes())?;
69    Ok(())
70}
71
72pub fn setup_file_struct() -> Result<()> {
73    if let Some(proj) = ProjectDirs::from(PROJECT.qualifier, PROJECT.org, PROJECT.app) {
74        if !proj.data_dir().exists() {
75            create_dir_all(proj.data_dir())?;
76        }
77        if !proj.config_dir().exists() {
78            create_dir_all(proj.config_dir())?;
79        }
80        let conf_file = format!("{}/{}", proj.config_dir().to_str().unwrap(), CONF_FILE);
81        if !Path::new(&conf_file).exists() {
82            println!("Configuration not detected: initiating config setup");
83            let color = get_confirm("Enable color", Some(true), false)?;
84            let mut host: String;
85            let mut port: u16;
86            loop {
87                host = get_input(
88                    "Enter Ollama server address",
89                    Some("localhost".to_owned()),
90                    color,
91                )?;
92                port = match color {
93                    true => Input::with_theme(&ColorfulTheme::default())
94                        .with_prompt("Enter Ollama server port")
95                        .default("11434".to_owned())
96                        .validate_with(|input: &String| -> Result<(), String> {
97                            validate_port_str(input)
98                        })
99                        .report(true)
100                        .interact_text()?,
101                    false => Input::new()
102                        .with_prompt("Enter Ollama server port")
103                        .default("11434".to_owned())
104                        .validate_with(|input: &String| -> Result<(), String> {
105                            validate_port_str(input)
106                        })
107                        .report(true)
108                        .interact_text()?,
109                }
110                .parse::<u16>()?;
111                let url = format!("http://{}:{}", host, port);
112                if reqwest::blocking::get(&url).is_ok() {
113                    break;
114                }
115                let msg = format!("Ollama server not found at {}", url);
116                fmt_print(&msg, ContentType::Error, color);
117            }
118            let model = get_input("Enter model", Some("llama3:latest".to_owned()), color)?;
119            let verbose = get_confirm("Enable verbose output", Some(true), color)?;
120            let save = get_confirm("Enable autosave on exit", Some(true), color)?;
121            let conf = Config {
122                host,
123                port,
124                model,
125                verbose,
126                color,
127                save,
128            };
129            let mut file = File::create(conf_file)?;
130            file.write_all(to_string(&conf)?.as_bytes())?;
131            fmt_print(
132                "NOTE: Params can be changed in config file.",
133                ContentType::Info,
134                color,
135            );
136        }
137        return Ok(());
138    }
139    Err(anyhow!("Could not create project directory"))
140}
141
142pub fn get_input(prompt: &str, default_opt: Option<String>, color: bool) -> Result<String> {
143    let (default, show_default) = match default_opt {
144        Some(s) => (s, true),
145        None => (String::new(), false),
146    };
147    let user_input: String = match color {
148        true => Input::with_theme(&ColorfulTheme::default())
149            .with_prompt(prompt)
150            .default(default)
151            .show_default(show_default)
152            .report(true)
153            .interact_text()?,
154        false => Input::new()
155            .with_prompt(prompt)
156            .default(default)
157            .show_default(show_default)
158            .report(true)
159            .interact_text()?,
160    };
161    Ok(user_input)
162}
163
164pub fn get_confirm(prompt: &str, default_opt: Option<bool>, color: bool) -> Result<bool> {
165    let (default, show_default) = match default_opt {
166        Some(b) => (b, true),
167        None => (false, false),
168    };
169    let ans = match color {
170        true => Confirm::with_theme(&ColorfulTheme::default())
171            .with_prompt(prompt)
172            .default(default)
173            .show_default(show_default)
174            .wait_for_newline(true)
175            .interact()?,
176        false => Confirm::new()
177            .with_prompt(prompt)
178            .default(default)
179            .show_default(show_default)
180            .wait_for_newline(true)
181            .interact()?,
182    };
183    Ok(ans)
184}
185
186pub fn get_config(matches: ArgMatches) -> Result<Config> {
187    let toml_string = read_file(&get_project_file(ProjFiles::Conf)?)?;
188    let mut conf: Config = toml::from_str(&toml_string)?;
189    if matches.value_source("host").is_some() {
190        conf.host = matches.get_one::<String>("host").unwrap().to_string();
191    }
192    if matches.value_source("port").is_some() {
193        conf.port = matches
194            .get_one::<String>("port")
195            .unwrap()
196            .to_string()
197            .parse::<u16>()?;
198    }
199    if matches.value_source("model").is_some() {
200        conf.model = matches.get_one::<String>("model").unwrap().to_string();
201    }
202    if matches.get_flag("verbose") {
203        conf.verbose = true;
204    }
205    if matches.get_flag("color") {
206        conf.color = true;
207    }
208    if matches.get_flag("save") {
209        conf.save = true;
210    }
211    ensure!(conf.port < 65535, "Port out of bounds");
212    let msg = format!(
213        "Ollama host {}:{} with model \"{}\"",
214        &conf.host, &conf.port, &conf.model
215    );
216    log(LogLevel::Info, "conf", &msg)?;
217    Ok(conf)
218}
219
220pub fn fmt_print(s: &str, content_type: ContentType, color: bool) {
221    if color {
222        match content_type {
223            ContentType::Error => eprintln!("{}", s.red()),
224            ContentType::Info => println!("{}", s.yellow().italic()),
225            ContentType::Answer => {
226                PrettyPrinter::new()
227                    .input_from_bytes(s.as_bytes())
228                    .grid(true)
229                    .language("markdown")
230                    .theme("DarkNeon")
231                    .print()
232                    .unwrap();
233            }
234            ContentType::Exit => println!("{}", s.green()),
235        }
236    } else {
237        match content_type {
238            ContentType::Error => eprintln!("{}", s),
239            _ => println!("{}", s),
240        }
241    }
242}
243
244pub fn get_project_file(file: ProjFiles) -> Result<String> {
245    if let Some(proj) = ProjectDirs::from(PROJECT.qualifier, PROJECT.org, PROJECT.app) {
246        match file {
247            ProjFiles::Conf => {
248                return Ok(format!(
249                    "{}/{}",
250                    proj.config_dir().to_str().unwrap(),
251                    CONF_FILE
252                ));
253            }
254            ProjFiles::Log => {
255                return Ok(format!(
256                    "{}/{}",
257                    proj.data_dir().to_str().unwrap(),
258                    LOG_FILE
259                ));
260            }
261            ProjFiles::Data => {
262                return Ok(format!("{}/{}", proj.data_dir().to_str().unwrap(), DB_FILE));
263            }
264        }
265    }
266    Err(anyhow!("Could not get project file"))
267}
268
269fn read_file(path: &str) -> Result<String> {
270    let mut s = String::new();
271    let mut f = File::open(path)?;
272    f.read_to_string(&mut s)?;
273    Ok(s)
274}
275
276fn validate_port_str(port_str: &str) -> Result<(), String> {
277    if port_str.parse::<u16>().is_err() {
278        return Err("Invalid port".to_owned());
279    }
280    Ok(())
281}
lib/lib.rs

lib/
lib.rs
