diff --git a/Cargo.lock b/Cargo.lock index 2a6ad2168..968a29880 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4231,6 +4231,7 @@ dependencies = [ "prometheus-client", "rayon", "redis 0.25.4", + "secrecy", "serde", "serde_json", "sqlx", diff --git a/Cargo.toml b/Cargo.toml index 69941a014..2e762414b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -241,7 +241,6 @@ prometheus-client = "0.22.0" brotli = "3.4.0" collab-stream = { path = "libs/collab-stream" } dotenvy = "0.15.7" -secrecy = { version = "0.8", features = ["serde"] } serde_json = "1.0.111" serde_repr = "0.1.18" serde = { version = "1.0.195", features = ["derive"] } @@ -292,7 +291,7 @@ chrono = { version = "0.4.39", features = [ ], default-features = false } http = "0.2.12" tokio-tungstenite = "0.20" - +secrecy = { version = "0.8.0", features = ["serde"] } # collaboration yrs = { version = "0.21.3", features = ["sync"] } collab = { version = "0.2.0" } diff --git a/libs/indexer/Cargo.toml b/libs/indexer/Cargo.toml index 2602a9f44..8b051934b 100644 --- a/libs/indexer/Cargo.toml +++ b/libs/indexer/Cargo.toml @@ -37,4 +37,5 @@ redis = { workspace = true, features = [ "tokio-comp", "connection-manager", ] } -tokio-util = "0.7.12" \ No newline at end of file +tokio-util = "0.7.12" +secrecy = { workspace = true, features = ["serde"] } diff --git a/libs/indexer/src/scheduler.rs b/libs/indexer/src/scheduler.rs index 7c014f0ad..c88163964 100644 --- a/libs/indexer/src/scheduler.rs +++ b/libs/indexer/src/scheduler.rs @@ -18,6 +18,7 @@ use database_entity::dto::AFCollabEmbeddedChunk; use infra::env_util::get_env_var; use rayon::prelude::*; use redis::aio::ConnectionManager; +use secrecy::{ExposeSecret, Secret}; use serde::{Deserialize, Serialize}; use sqlx::PgPool; use std::cmp::max; @@ -49,7 +50,7 @@ pub struct IndexerScheduler { #[derive(Debug)] pub struct IndexerConfiguration { pub enable: bool, - pub openai_api_key: String, + pub openai_api_key: Secret, /// High watermark for the number of embeddings that can be buffered before being written to the database. pub embedding_buffer_size: usize, } @@ -129,7 +130,7 @@ impl IndexerScheduler { } // if openai api key is empty, return false - if self.config.openai_api_key.is_empty() { + if self.config.openai_api_key.expose_secret().is_empty() { return false; } @@ -141,14 +142,14 @@ impl IndexerScheduler { } pub(crate) fn create_embedder(&self) -> Result { - if self.config.openai_api_key.is_empty() { + if self.config.openai_api_key.expose_secret().is_empty() { return Err(AppError::AIServiceUnavailable( "OpenAI API key is empty".to_string(), )); } Ok(Embedder::OpenAI(open_ai::Embedder::new( - self.config.openai_api_key.clone(), + self.config.openai_api_key.expose_secret().clone(), ))) } diff --git a/services/appflowy-collaborate/src/application.rs b/services/appflowy-collaborate/src/application.rs index 8643600b0..7f1e3aa43 100644 --- a/services/appflowy-collaborate/src/application.rs +++ b/services/appflowy-collaborate/src/application.rs @@ -15,7 +15,7 @@ use aws_sdk_s3::operation::create_bucket::CreateBucketError; use aws_sdk_s3::types::{ BucketInfo, BucketLocationConstraint, BucketType, CreateBucketConfiguration, }; -use secrecy::ExposeSecret; +use secrecy::{ExposeSecret, Secret}; use sqlx::postgres::PgPoolOptions; use sqlx::PgPool; use tracing::info; @@ -158,7 +158,7 @@ pub async fn init_state(config: &Config, rt_cmd_tx: CLCommandSender) -> Result() .unwrap_or(true), - openai_api_key: get_env_var("APPFLOWY_AI_OPENAI_API_KEY", ""), + openai_api_key: Secret::new(get_env_var("APPFLOWY_AI_OPENAI_API_KEY", "")), embedding_buffer_size: get_env_var("APPFLOWY_INDEXER_EMBEDDING_BUFFER_SIZE", "2000") .parse::() .unwrap_or(2000), diff --git a/services/appflowy-worker/Cargo.toml b/services/appflowy-worker/Cargo.toml index a448000b0..4d5544ad8 100644 --- a/services/appflowy-worker/Cargo.toml +++ b/services/appflowy-worker/Cargo.toml @@ -46,7 +46,7 @@ sqlx = { workspace = true, default-features = false, features = [ "chrono", "migrate", ] } -secrecy = { version = "0.8", features = ["serde"] } +secrecy = { workspace = true, features = ["serde"] } aws-sdk-s3 = { version = "1.36.0", features = [ "behavior-version-latest", "rt-tokio", diff --git a/services/appflowy-worker/src/application.rs b/services/appflowy-worker/src/application.rs index af46212a9..187b053ba 100644 --- a/services/appflowy-worker/src/application.rs +++ b/services/appflowy-worker/src/application.rs @@ -11,7 +11,6 @@ use crate::import_worker::email_notifier::EmailNotifier; use crate::s3_client::S3ClientImpl; use axum::Router; -use secrecy::ExposeSecret; use crate::mailer::AFWorkerMailer; use crate::metric::ImportMetrics; @@ -24,6 +23,7 @@ use indexer::metrics::EmbeddingMetrics; use indexer::thread_pool::ThreadPoolNoAbortBuilder; use infra::env_util::get_env_var; use mailer::sender::Mailer; +use secrecy::{ExposeSecret, Secret}; use std::sync::{Arc, Once}; use std::time::Duration; use tokio::net::TcpListener; @@ -144,7 +144,10 @@ pub async fn create_app(listener: TcpListener, config: Config) -> Result<(), Err enable: appflowy_collaborate::config::get_env_var("APPFLOWY_INDEXER_ENABLED", "true") .parse::() .unwrap_or(true), - open_api_key: appflowy_collaborate::config::get_env_var("APPFLOWY_AI_OPENAI_API_KEY", ""), + open_api_key: Secret::new(appflowy_collaborate::config::get_env_var( + "APPFLOWY_AI_OPENAI_API_KEY", + "", + )), tick_interval_secs: 10, }, )); diff --git a/services/appflowy-worker/src/indexer_worker/worker.rs b/services/appflowy-worker/src/indexer_worker/worker.rs index fd0d20253..7e776e7c9 100644 --- a/services/appflowy-worker/src/indexer_worker/worker.rs +++ b/services/appflowy-worker/src/indexer_worker/worker.rs @@ -15,6 +15,7 @@ use indexer::vector::embedder::Embedder; use indexer::vector::open_ai; use rayon::prelude::*; use redis::aio::ConnectionManager; +use secrecy::{ExposeSecret, Secret}; use sqlx::PgPool; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -23,10 +24,9 @@ use tokio::sync::RwLock; use tokio::time::{interval, MissedTickBehavior}; use tracing::{error, info, trace}; -#[derive(Debug)] pub struct BackgroundIndexerConfig { pub enable: bool, - pub open_api_key: String, + pub open_api_key: Secret, pub tick_interval_secs: u64, } @@ -42,7 +42,7 @@ pub async fn run_background_indexer( return; } - if config.open_api_key.is_empty() { + if config.open_api_key.expose_secret().is_empty() { error!("OpenAI API key is not set. Stop background indexer"); return; } @@ -243,5 +243,7 @@ fn handle_task( } fn create_embedder(config: &BackgroundIndexerConfig) -> Embedder { - Embedder::OpenAI(open_ai::Embedder::new(config.open_api_key.clone())) + Embedder::OpenAI(open_ai::Embedder::new( + config.open_api_key.expose_secret().clone(), + )) } diff --git a/src/application.rs b/src/application.rs index 8f76e2ec5..7f96a97b3 100644 --- a/src/application.rs +++ b/src/application.rs @@ -327,7 +327,7 @@ pub async fn init_state(config: &Config, rt_cmd_tx: CLCommandSender) -> Result() .unwrap_or(true), - openai_api_key: get_env_var("APPFLOWY_AI_OPENAI_API_KEY", ""), + openai_api_key: Secret::new(get_env_var("APPFLOWY_AI_OPENAI_API_KEY", "")), embedding_buffer_size: appflowy_collaborate::config::get_env_var( "APPFLOWY_INDEXER_EMBEDDING_BUFFER_SIZE", "5000",