From 797d875755e1e2f464bf6874397f07dbd9c2086a Mon Sep 17 00:00:00 2001 From: echozyr2001 Date: Wed, 7 May 2025 20:44:10 +0800 Subject: [PATCH 1/2] Add Credentials struct for simplified client initialization --- src/openai.rs | 181 +++++++++++++++++++++++++++++++++++++++++++------- 1 file changed, 158 insertions(+), 23 deletions(-) diff --git a/src/openai.rs b/src/openai.rs index 0f24f11..a55f72d 100644 --- a/src/openai.rs +++ b/src/openai.rs @@ -1,57 +1,108 @@ use serde::{Deserialize, Serialize}; use ureq::{Agent, AgentBuilder}; -#[derive(Debug, Serialize, Deserialize)] +#[derive(Debug, Serialize, Deserialize, Clone)] pub struct Auth { pub api_key: String, pub organization: Option, } -impl Clone for Auth { - fn clone(&self) -> Self { - Self { api_key: self.api_key.clone(), organization: self.organization.clone() } - } -} - -#[allow(dead_code)] impl Auth { pub fn new(api_key: &str) -> Auth { Auth { api_key: api_key.to_string(), organization: None } } + pub fn with_organization(mut self, organization: &str) -> Self { + self.organization = Some(organization.to_string()); + self + } + pub fn from_env() -> Result { let api_key = std::env::var("OPENAI_API_KEY").map_err(|_| "Missing OPENAI_API_KEY".to_string())?; - Ok(Self { api_key, organization: None }) + let organization = std::env::var("OPENAI_ORGANIZATION").ok(); + + Ok(Self { api_key, organization }) } } -#[derive(Debug)] -pub struct OpenAI { +/// Container for API credentials and URL configuration +#[derive(Debug, Clone)] +pub struct Credentials { pub auth: Auth, pub api_url: String, - pub(crate) agent: Agent, } -impl Clone for OpenAI { - fn clone(&self) -> Self { - Self { auth: self.auth.clone(), api_url: self.api_url.clone(), agent: self.agent.clone() } +impl Credentials { + /// Create new credentials with the specified API key and URL + pub fn new(api_key: &str, api_url: &str) -> Self { + Self { auth: Auth::new(api_key), api_url: api_url.to_string() } + } + + /// Add organization ID to these credentials + pub fn with_organization(mut self, organization: &str) -> Self { + self.auth.organization = Some(organization.to_string()); + self + } + + /// Load credentials from environment variables: + /// - OPENAI_API_KEY: Required - your OpenAI API key + /// - OPENAI_API_URL: Optional - defaults to "https://api.openai.com/v1/" + /// - OPENAI_ORGANIZATION: Optional - your organization ID + pub fn from_env() -> Result { + let api_key = + std::env::var("OPENAI_API_KEY").map_err(|_| "Missing OPENAI_API_KEY".to_string())?; + + let api_url = std::env::var("OPENAI_API_URL") + .unwrap_or_else(|_| "https://api.openai.com/v1/".to_string()); + + Ok(Self { + auth: Auth { api_key, organization: std::env::var("OPENAI_ORGANIZATION").ok() }, + api_url, + }) } } -#[allow(dead_code)] +#[derive(Debug, Clone)] +pub struct OpenAI { + pub auth: Auth, + pub api_url: String, + pub(crate) agent: Agent, +} + impl OpenAI { - pub fn new(auth: Auth, api_url: &str) -> OpenAI { - OpenAI { auth, api_url: api_url.to_string(), agent: AgentBuilder::new().build() } + pub fn new(auth: Auth, api_url: &str) -> Self { + Self { auth, api_url: api_url.to_string(), agent: AgentBuilder::new().build() } } - pub fn set_proxy(mut self, proxy: &str) -> OpenAI { - let proxy = ureq::Proxy::new(proxy).unwrap(); + /// Initialize client from environment variables + /// A convenient shorthand for Credentials::from_env() + from_credentials() + pub fn from_env() -> Result { + let credentials = Credentials::from_env()?; + Ok(Self::from_credentials(credentials)) + } + + /// Initialize client using a Credentials object + pub fn from_credentials(credentials: Credentials) -> Self { + Self { + auth: credentials.auth, + api_url: credentials.api_url, + agent: AgentBuilder::new().build(), + } + } + + pub fn builder() -> OpenAIBuilder { + OpenAIBuilder::new() + } + + pub fn set_proxy(mut self, proxy: &str) -> Self { + let proxy = ureq::Proxy::new(proxy).map_err(|e| format!("Invalid proxy: {}", e)).unwrap(); + self.agent = ureq::AgentBuilder::new().proxy(proxy).build(); self } - pub fn use_env_proxy(mut self) -> OpenAI { + pub fn use_env_proxy(mut self) -> Self { let proxy = match (std::env::var("http_proxy"), std::env::var("https_proxy")) { (Ok(http_proxy), _) => Some(http_proxy), (_, Ok(https_proxy)) => Some(https_proxy), @@ -60,12 +111,90 @@ impl OpenAI { None }, }; + if let Some(proxy) = proxy { - let proxy = ureq::Proxy::new(&proxy).unwrap(); - self.agent = ureq::AgentBuilder::new().proxy(proxy).build(); + if let Ok(proxy) = ureq::Proxy::new(&proxy) { + self.agent = ureq::AgentBuilder::new().proxy(proxy).build(); + } + } + self + } +} + +pub struct OpenAIBuilder { + api_key: Option, + api_url: Option, + organization: Option, + use_proxy: bool, + proxy: Option, +} + +impl OpenAIBuilder { + fn new() -> Self { + Self { api_key: None, api_url: None, organization: None, use_proxy: false, proxy: None } + } + + pub fn api_key(mut self, api_key: &str) -> Self { + self.api_key = Some(api_key.to_string()); + self + } + + pub fn api_url(mut self, api_url: &str) -> Self { + self.api_url = Some(api_url.to_string()); + self + } + + pub fn organization(mut self, organization: &str) -> Self { + self.organization = Some(organization.to_string()); + self + } + + pub fn use_env_proxy(mut self) -> Self { + self.use_proxy = true; + self + } + + pub fn proxy(mut self, proxy: &str) -> Self { + self.proxy = Some(proxy.to_string()); + self + } + + pub fn from_env(mut self) -> Self { + if self.api_key.is_none() { + self.api_key = std::env::var("OPENAI_API_KEY").ok(); + } + + if self.api_url.is_none() { + self.api_url = std::env::var("OPENAI_API_URL").ok(); + } + + if self.organization.is_none() { + self.organization = std::env::var("OPENAI_ORGANIZATION").ok(); } + self } + + pub fn build(self) -> Result { + let api_key = self.api_key.ok_or_else(|| "API key is required".to_string())?; + + let api_url = self.api_url.unwrap_or_else(|| "https://api.openai.com/v1/".to_string()); + + let mut auth = Auth::new(&api_key); + if let Some(org) = self.organization { + auth.organization = Some(org); + } + + let mut client = OpenAI::new(auth, &api_url); + + if let Some(proxy) = self.proxy { + client = client.set_proxy(&proxy); + } else if self.use_proxy { + client = client.use_env_proxy(); + } + + Ok(client) + } } #[cfg(test)] @@ -73,3 +202,9 @@ pub fn new_test_openai() -> OpenAI { let auth = Auth::from_env().unwrap(); OpenAI::new(auth, "https://api.openai.com/v1/").use_env_proxy() } + +#[cfg(test)] +pub fn new_test_openai_with_credentials() -> OpenAI { + let credentials = Credentials::from_env().unwrap(); + OpenAI::from_credentials(credentials).use_env_proxy() +} From a2c2460dc0ad42124b835f53b6941beccab93d84 Mon Sep 17 00:00:00 2001 From: echozyr2001 Date: Wed, 7 May 2025 20:48:34 +0800 Subject: [PATCH 2/2] Add test for Credentials struct --- tests/test_credentials.rs | 77 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 77 insertions(+) create mode 100644 tests/test_credentials.rs diff --git a/tests/test_credentials.rs b/tests/test_credentials.rs new file mode 100644 index 0000000..eb0eb9a --- /dev/null +++ b/tests/test_credentials.rs @@ -0,0 +1,77 @@ +use openai_api_rust::openai::{Credentials, OpenAI}; +use std::env; + +#[test] +fn test_credentials_from_env() { + // Save original env vars to restore them after test + let original_api_key = env::var("OPENAI_API_KEY").ok(); + let original_api_url = env::var("OPENAI_API_URL").ok(); + let original_organization = env::var("OPENAI_ORGANIZATION").ok(); + + // Set test values + env::set_var("OPENAI_API_KEY", "test_key"); + env::set_var("OPENAI_API_URL", "https://test.api.com/v1/"); + env::set_var("OPENAI_ORGANIZATION", "test_org"); + + // Test credentials creation + let credentials = Credentials::from_env().unwrap(); + assert_eq!(credentials.auth.api_key, "test_key"); + assert_eq!(credentials.api_url, "https://test.api.com/v1/"); + assert_eq!(credentials.auth.organization, Some("test_org".to_string())); + + // Test OpenAI client creation with credentials + let client = OpenAI::from_credentials(credentials); + assert_eq!(client.auth.api_key, "test_key"); + assert_eq!(client.api_url, "https://test.api.com/v1/"); + assert_eq!(client.auth.organization, Some("test_org".to_string())); + + // Test direct creation from env + let client = OpenAI::from_env().unwrap(); + assert_eq!(client.auth.api_key, "test_key"); + assert_eq!(client.api_url, "https://test.api.com/v1/"); + assert_eq!(client.auth.organization, Some("test_org".to_string())); + + // Test with default URL + env::remove_var("OPENAI_API_URL"); + let credentials = Credentials::from_env().unwrap(); + assert_eq!(credentials.api_url, "https://api.openai.com/v1/"); + + // Restore original env vars + match original_api_key { + Some(val) => env::set_var("OPENAI_API_KEY", val), + None => env::remove_var("OPENAI_API_KEY"), + } + + match original_api_url { + Some(val) => env::set_var("OPENAI_API_URL", val), + None => env::remove_var("OPENAI_API_URL"), + } + + match original_organization { + Some(val) => env::set_var("OPENAI_ORGANIZATION", val), + None => env::remove_var("OPENAI_ORGANIZATION"), + } +} + +#[test] +fn test_credentials_builder() { + // Create credentials directly + let credentials = Credentials::new("manual_key", "https://manual.api.com/v1/") + .with_organization("manual_org"); + + assert_eq!(credentials.auth.api_key, "manual_key"); + assert_eq!(credentials.api_url, "https://manual.api.com/v1/"); + assert_eq!(credentials.auth.organization, Some("manual_org".to_string())); + + // Create OpenAI client using the builder pattern + let client = OpenAI::builder() + .api_key("builder_key") + .api_url("https://builder.api.com/v1/") + .organization("builder_org") + .build() + .unwrap(); + + assert_eq!(client.auth.api_key, "builder_key"); + assert_eq!(client.api_url, "https://builder.api.com/v1/"); + assert_eq!(client.auth.organization, Some("builder_org".to_string())); +}