core embedder/llm DATA_SOURCES api keys (#6496)

dust-tt · Jul 24, 2024 · a0fa5b9 · a0fa5b9
1 parent ded2575
commit a0fa5b9
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 48 deletions.
diff --git a/core/src/providers/mistral.rs b/core/src/providers/mistral.rs
@@ -882,24 +882,20 @@ impl LLM for MistralAILLM {
         self.id.clone()
     }
     async fn initialize(&mut self, credentials: Credentials) -> Result<()> {
-        match std::env::var("CORE_DATA_SOURCES_MISTRAL_API_KEY") {
-            Ok(key) => {
-                self.api_key = Some(key);
+        match credentials.get("MISTRAL_API_KEY") {
+            Some(api_key) => {
+                self.api_key = Some(api_key.clone());
             }
-            Err(_) => match credentials.get("MISTRAL_API_KEY") {
-                Some(api_key) => {
-                    self.api_key = Some(api_key.clone());
+            None => match tokio::task::spawn_blocking(|| std::env::var("MISTRAL_API_KEY")).await? {
+                Ok(key) => {
+                    self.api_key = Some(key);
                 }
-                None => match std::env::var("MISTRAL_API_KEY") {
-                    Ok(key) => {
-                        self.api_key = Some(key);
-                    }
-                    Err(_) => Err(anyhow!(
-                        "Credentials or environment variable `MISTRAL_API_KEY` is not set."
-                    ))?,
-                },
+                Err(_) => Err(anyhow!(
+                    "Credentials or environment variable `MISTRAL_API_KEY` is not set."
+                ))?,
             },
         }
+
         Ok(())
     }
 
@@ -1109,20 +1105,24 @@ impl Embedder for MistralEmbedder {
             ));
         }
 
-        match credentials.get("MISTRAL_API_KEY") {
-            Some(api_key) => {
-                self.api_key = Some(api_key.clone());
+        match std::env::var("CORE_DATA_SOURCES_MISTRAL_API_KEY") {
+            Ok(key) => {
+                self.api_key = Some(key);
             }
-            None => match tokio::task::spawn_blocking(|| std::env::var("MISTRAL_API_KEY")).await? {
-                Ok(key) => {
-                    self.api_key = Some(key);
+            Err(_) => match credentials.get("MISTRAL_API_KEY") {
+                Some(api_key) => {
+                    self.api_key = Some(api_key.clone());
                 }
-                Err(_) => Err(anyhow!(
-                    "Credentials or environment variable `MISTRAL_API_KEY` is not set."
-                ))?,
+                None => match std::env::var("MISTRAL_API_KEY") {
+                    Ok(key) => {
+                        self.api_key = Some(key);
+                    }
+                    Err(_) => Err(anyhow!(
+                        "Credentials or environment variable `MISTRAL_API_KEY` is not set."
+                    ))?,
+                },
             },
         }
-
         Ok(())
     }
 

diff --git a/core/src/providers/openai.rs b/core/src/providers/openai.rs
@@ -1721,23 +1721,17 @@ impl LLM for OpenAILLM {
     }
 
     async fn initialize(&mut self, credentials: Credentials) -> Result<()> {
-        // Give priority to `CORE_DATA_SOURCES_OPENAI_API_KEY` env variable
-        match std::env::var("CORE_DATA_SOURCES_OPENAI_API_KEY") {
-            Ok(key) => {
-                self.api_key = Some(key);
+        match credentials.get("OPENAI_API_KEY") {
+            Some(api_key) => {
+                self.api_key = Some(api_key.clone());
             }
-            Err(_) => match credentials.get("OPENAI_API_KEY") {
-                Some(api_key) => {
-                    self.api_key = Some(api_key.clone());
+            None => match tokio::task::spawn_blocking(|| std::env::var("OPENAI_API_KEY")).await? {
+                Ok(key) => {
+                    self.api_key = Some(key);
                 }
-                None => match std::env::var("OPENAI_API_KEY") {
-                    Ok(key) => {
-                        self.api_key = Some(key);
-                    }
-                    Err(_) => Err(anyhow!(
-                        "Credentials or environment variable `OPENAI_API_KEY` is not set."
-                    ))?,
-                },
+                Err(_) => Err(anyhow!(
+                    "Credentials or environment variable `OPENAI_API_KEY` is not set."
+                ))?,
             },
         }
         Ok(())
@@ -2154,17 +2148,23 @@ impl Embedder for OpenAIEmbedder {
             ));
         }
 
-        match credentials.get("OPENAI_API_KEY") {
-            Some(api_key) => {
-                self.api_key = Some(api_key.clone());
+        // Give priority to `CORE_DATA_SOURCES_OPENAI_API_KEY` env variable
+        match std::env::var("CORE_DATA_SOURCES_OPENAI_API_KEY") {
+            Ok(key) => {
+                self.api_key = Some(key);
             }
-            None => match tokio::task::spawn_blocking(|| std::env::var("OPENAI_API_KEY")).await? {
-                Ok(key) => {
-                    self.api_key = Some(key);
+            Err(_) => match credentials.get("OPENAI_API_KEY") {
+                Some(api_key) => {
+                    self.api_key = Some(api_key.clone());
                 }
-                Err(_) => Err(anyhow!(
-                    "Credentials or environment variable `OPENAI_API_KEY` is not set."
-                ))?,
+                None => match std::env::var("OPENAI_API_KEY") {
+                    Ok(key) => {
+                        self.api_key = Some(key);
+                    }
+                    Err(_) => Err(anyhow!(
+                        "Credentials or environment variable `OPENAI_API_KEY` is not set."
+                    ))?,
+                },
             },
         }
         Ok(())