openrouter_rs/command/
model.rs

1use std::{
2    fs::{self, File},
3    io::{BufRead, BufReader, BufWriter, Error, ErrorKind, Write},
4};
5
6use dirs::config_dir;
7
8use crate::error::OpenRouterError;
9
10/// Stores the user's preferred model names in a file within the user's config directory.
11/// This allows the application to remember the user's model preferences across sessions.
12///
13/// # Arguments
14///
15/// * `model_names` - A slice of model names to set as preferred.
16///
17/// # Returns
18///
19/// * `Result<(), OpenRouterError>` - Ok if the operation was successful, or an error if it failed.
20pub fn set_preferred_models(model_names: &[&str]) -> Result<(), OpenRouterError> {
21    // Retrieve the user's config directory to store the preferred models
22    let config_dir = config_dir()
23        .ok_or_else(|| Error::new(ErrorKind::NotFound, "Config directory not found"))?;
24
25    // Ensure the openrouter directory exists within the config directory
26    let openrouter_dir = config_dir.join("openrouter");
27    if !openrouter_dir.exists() {
28        fs::create_dir_all(&openrouter_dir)?;
29    }
30
31    // Create or open the selected_models file within the openrouter directory
32    let selected_models_file = openrouter_dir.join("selected_models");
33    let mut file = BufWriter::new(File::create(selected_models_file)?);
34
35    // Write each model name to the file on a new line
36    for model_name in model_names {
37        writeln!(file, "{}", model_name)?;
38    }
39    Ok(())
40}
41
42/// Retrieves the user's preferred model names from a file within the user's config directory.
43/// This allows the application to load the user's model preferences when it starts.
44///
45/// # Returns
46///
47/// * `Result<Vec<String>, OpenRouterError>` - A vector of preferred model names if found, or an error if it failed.
48pub fn get_preferred_models() -> Result<Vec<String>, OpenRouterError> {
49    // Retrieve the user's config directory to read the preferred models
50    let config_dir = config_dir()
51        .ok_or_else(|| Error::new(ErrorKind::NotFound, "Config directory not found"))?;
52
53    // Ensure the openrouter directory exists within the config directory
54    let openrouter_dir = config_dir.join("openrouter");
55    if !openrouter_dir.exists() {
56        fs::create_dir_all(&openrouter_dir)?;
57    }
58
59    // Open the selected_models file within the openrouter directory
60    let selected_models_file = openrouter_dir.join("selected_models");
61    let file = BufReader::new(File::open(selected_models_file)?);
62
63    // Read each line from the file and collect them into a vector of model names
64    let mut model_names = Vec::new();
65    for line in file.lines() {
66        let line = line?;
67        model_names.push(line);
68    }
69    Ok(model_names)
70}