From 97e6a058fef945580953ed42f99ac4f5434416cf Mon Sep 17 00:00:00 2001
From: theoforger <theoforger@proton.me>
Date: Wed, 18 Sep 2024 08:22:31 -0400
Subject: [PATCH] Move all environment variables handling inside Instance

---
 src/api/mod.rs | 15 +++++++++------
 src/main.rs    | 17 ++++-------------
 2 files changed, 13 insertions(+), 19 deletions(-)

diff --git a/src/api/mod.rs b/src/api/mod.rs
index 3ff7e16..b95b552 100644
--- a/src/api/mod.rs
+++ b/src/api/mod.rs
@@ -16,25 +16,28 @@ impl Instance {
     pub fn new() -> Result<Self, Box<dyn std::error::Error>> {
         dotenv().ok();
 
-        let base_url = env::var("OPENAI_API_BASE_URL")
-            .map_err(|_| "Cannot read environment variable: OPENAI_API_BASE_URL")?;
-
+        let base_url = Self::get_env_var("OPENAI_API_BASE_URL")?;
         let base_url = if !base_url.ends_with('/') {
             format!("{}/", base_url)
         } else {
             base_url
         };
-
-        let key = env::var("API_KEY").map_err(|_| "Cannot read environment variable: API_KEY")?;
+        let key = Self::get_env_var("API_KEY")?;
+        let model_id = Self::get_env_var("DEFAULT_MODEL_ID")?;
 
         Ok(Self {
             client: reqwest::Client::new(),
             base_url,
             key,
-            model_id: "".to_string(),
+            model_id,
         })
     }
 
+    fn get_env_var(var_name: &str) -> Result<String, Box<dyn std::error::Error>> {
+        env::var(var_name)
+            .map_err(|_| format!("Cannot read environment variable: {}", var_name).into())
+    }
+
     pub async fn set_model_id(
         &mut self,
         model_id: String,
diff --git a/src/main.rs b/src/main.rs
index 6b88eb0..838ffc5 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,7 +1,4 @@
-use std::env;
-
 use clap::Parser;
-use dotenv::dotenv;
 
 use mastermind::api::Instance;
 use mastermind::*;
@@ -10,7 +7,6 @@ use mastermind::*;
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
     // Read arguments and environment variables
     let args = Args::parse();
-    dotenv().ok();
 
     // Create an API instance
     let mut api_instance = Instance::new()?;
@@ -22,14 +18,9 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     }
 
     // If -m is set, use a preferred language model
-    // Otherwise, set the default
-    let model_id = match args.model {
-        Some(id) => id,
-        None => env::var("DEFAULT_MODEL_ID").map_err(|_| {
-            "Could not read environment variable: DEFAULT_MODEL_ID. Use -m to specify a language model"
-        })?,
-    };
-    api_instance.set_model_id(model_id).await?;
+    if let Some(model_id) = args.model {
+        api_instance.set_model_id(model_id).await?;
+    }
 
     // Attempt to read words from the two files
     let link_words = read_words_from_file(args.to_link.unwrap()).map_err(|e| e.to_string())?;
@@ -44,7 +35,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     if clue_collection.is_empty() {
         println!("The language model didn't return any useful clues. Maybe try again?");
     } else if let Some(output_path) = args.output {
-        println!("-o option is present, writing to file...");
+        println!("Writing to file '{}'...", output_path.display());
         write_content_to_file(output_path, clue_collection.output())?;
     } else {
         clue_collection.display();