a/
lib.rs

1use llm::LLMClient;
2use std::error::Error;
3use std::io;
4
5pub mod llm;
6pub mod util;
7
8pub use util::*;
9
10pub fn gather_args(args: &mut Vec<String>) -> Result<(String, String), Box<dyn Error>> {
11    let lang;
12    let mut prompt = String::new();
13
14    if args.is_empty() {
15        return Err("Received zero arguments".into());
16    }
17
18    args.remove(0);
19    if args.is_empty() {
20        if let Err(e) = io::stdin().read_line(&mut prompt) {
21            println!("Could not read from stdin: {e}");
22            return Err(e.into());
23        }
24
25        let words: Vec<String> = prompt.split_whitespace().map(str::to_string).collect();
26        if words.is_empty() {
27            return Err("Please supply a prompt".into());
28        }
29
30        if words[0] != "a" {
31            lang = words[0].to_string();
32        } else if words.len() >= 2 {
33            lang = words[1].to_string();
34        } else {
35            return Err("Please supply a prompt".into());
36        }
37    } else {
38        lang = args[0].clone();
39        prompt = args.join(" ");
40    }
41
42    Ok((prompt, lang))
43}
44
45pub fn prompt(
46    prompt: &str,
47    model: &str,
48    api_key: &str,
49    base_url: &str,
50) -> Result<String, Box<dyn Error>> {
51    let client = LLMClient::new(model, base_url, api_key);
52    let mut response = client.prompt(prompt.to_string())?;
53
54    while response.starts_with('\n') {
55        response.remove(0);
56    }
57    response.push('\n');
58
59    Ok(response)
60}