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}