From 63ba9703aecc49bcbf508e99c146343e5e96a95b Mon Sep 17 00:00:00 2001 From: erhant Date: Tue, 15 Oct 2024 00:28:59 +0300 Subject: [PATCH] fix env spacing and `"` issues --- Cargo.lock | 6 +++--- Cargo.toml | 2 +- workflows/src/apis/jina.rs | 4 +++- workflows/src/apis/serper.rs | 5 ++++- workflows/src/providers/openai.rs | 5 ++++- workflows/src/utils.rs | 24 ++++++++++++++++++++++-- 6 files changed, 37 insertions(+), 9 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index e1f0468..b255acf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -980,7 +980,7 @@ dependencies = [ [[package]] name = "dkn-compute" -version = "0.2.15" +version = "0.2.16" dependencies = [ "async-trait", "base64 0.22.1", @@ -1012,7 +1012,7 @@ dependencies = [ [[package]] name = "dkn-p2p" -version = "0.2.15" +version = "0.2.16" dependencies = [ "env_logger 0.11.5", "eyre", @@ -1024,7 +1024,7 @@ dependencies = [ [[package]] name = "dkn-workflows" -version = "0.2.15" +version = "0.2.16" dependencies = [ "dotenvy", "env_logger 0.11.5", diff --git a/Cargo.toml b/Cargo.toml index 4c0c9b5..5a54e10 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -7,7 +7,7 @@ default-members = ["compute"] [workspace.package] edition = "2021" -version = "0.2.15" +version = "0.2.16" license = "Apache-2.0" readme = "README.md" diff --git a/workflows/src/apis/jina.rs b/workflows/src/apis/jina.rs index 7879fd8..3ff0a71 100644 --- a/workflows/src/apis/jina.rs +++ b/workflows/src/apis/jina.rs @@ -2,6 +2,8 @@ use eyre::{eyre, Context, Result}; use reqwest::Client; use std::env; +use crate::utils::safe_read_env; + /// Makes a request for `example.com`. const JINA_EXAMPLE_ENDPOINT: &str = "https://r.jina.ai/https://example.com"; const ENV_VAR_NAME: &str = "JINA_API_KEY"; @@ -17,7 +19,7 @@ impl JinaConfig { /// Looks at the environment variables for Jina API key. pub fn new() -> Self { Self { - api_key: env::var(ENV_VAR_NAME).ok(), + api_key: safe_read_env(env::var(ENV_VAR_NAME)), } } diff --git a/workflows/src/apis/serper.rs b/workflows/src/apis/serper.rs index 56290de..0d916ec 100644 --- a/workflows/src/apis/serper.rs +++ b/workflows/src/apis/serper.rs @@ -2,6 +2,8 @@ use eyre::{eyre, Context, Result}; use reqwest::Client; use std::env; +use crate::utils::safe_read_env; + /// Makes a search request. const SERPER_EXAMPLE_ENDPOINT: &str = "https://google.serper.dev/search"; const ENV_VAR_NAME: &str = "SERPER_API_KEY"; @@ -17,7 +19,7 @@ impl SerperConfig { /// Looks at the environment variables for Serper API key. pub fn new() -> Self { Self { - api_key: env::var(ENV_VAR_NAME).ok(), + api_key: safe_read_env(env::var(ENV_VAR_NAME)), } } @@ -46,6 +48,7 @@ impl SerperConfig { log::debug!("Serper API key not found, skipping Serper check"); return Ok(()); }; + println!("API KEY: {}", api_key); log::info!("Serper API key found, checking Serper service"); // make a dummy request diff --git a/workflows/src/providers/openai.rs b/workflows/src/providers/openai.rs index 220b6a3..c61390a 100644 --- a/workflows/src/providers/openai.rs +++ b/workflows/src/providers/openai.rs @@ -3,7 +3,10 @@ use ollama_workflows::Model; use reqwest::Client; use serde::Deserialize; +use crate::utils::safe_read_env; + const OPENAI_MODELS_API: &str = "https://api.openai.com/v1/models"; +const ENV_VAR_NAME: &str = "OPENAI_API_KEY"; /// [Model](https://platform.openai.com/docs/api-reference/models/object) API object. #[derive(Debug, Clone, Deserialize)] @@ -39,7 +42,7 @@ impl OpenAIConfig { /// Looks at the environment variables for OpenAI API key. pub fn new() -> Self { Self { - api_key: std::env::var("OPENAI_API_KEY").ok(), + api_key: safe_read_env(std::env::var(ENV_VAR_NAME)), } } diff --git a/workflows/src/utils.rs b/workflows/src/utils.rs index 030f0e2..2e7b800 100644 --- a/workflows/src/utils.rs +++ b/workflows/src/utils.rs @@ -17,12 +17,20 @@ pub fn split_csv_line(input: &str) -> Vec { .collect::>() } +/// Reads an environment variable and trims whitespace and `"` from both ends. +/// If the trimmed value is empty, returns `None`. +pub fn safe_read_env(var: Result) -> Option { + var.map(|s| s.trim_matches('"').trim().to_string()) + .ok() + .filter(|s| !s.is_empty()) +} + #[cfg(test)] mod tests { use super::*; #[test] - fn test_example() { + fn test_abc_csv() { // should ignore whitespaces and `"` at both ends, and ignore empty items let input = "\"a, b , c ,, \""; let expected = vec!["a".to_string(), "b".to_string(), "c".to_string()]; @@ -30,7 +38,19 @@ mod tests { } #[test] - fn test_empty() { + fn test_empty_csv() { assert!(split_csv_line(Default::default()).is_empty()); } + + #[test] + fn test_var_read() { + let var = Ok("\" value \"".to_string()); + assert_eq!(safe_read_env(var), Some("value".to_string())); + + let var = Ok("\" \"".to_string()); + assert!(safe_read_env(var).is_none()); + + let var = Err(std::env::VarError::NotPresent); + assert!(safe_read_env(var).is_none()); + } }