From 0e50650aa8aa95d19f0f4d661a1005c5d065d3bb Mon Sep 17 00:00:00 2001 From: Bartosz Sypytkowski Date: Tue, 18 Jun 2024 10:51:11 +0200 Subject: [PATCH] chore: use appflowy ai client for embeddings in search api --- .github/workflows/integration_test.yml | 2 -- Cargo.lock | 37 ----------------------- Cargo.toml | 2 -- docker-compose.yml | 2 +- src/api/search.rs | 23 ++++++-------- src/application.rs | 11 ------- src/biz/search/ops.rs | 42 +++++++++++--------------- src/config/config.rs | 4 --- src/state.rs | 1 - 9 files changed, 28 insertions(+), 96 deletions(-) diff --git a/.github/workflows/integration_test.yml b/.github/workflows/integration_test.yml index b91a84dd..fe3f83a7 100644 --- a/.github/workflows/integration_test.yml +++ b/.github/workflows/integration_test.yml @@ -94,8 +94,6 @@ jobs: sed -i 's/APPFLOWY_MAILER_SMTP_USERNAME=.*/APPFLOWY_MAILER_SMTP_USERNAME=${{ secrets.CI_GOTRUE_SMTP_USER }}/' .env sed -i 's/APPFLOWY_MAILER_SMTP_PASSWORD=.*/APPFLOWY_MAILER_SMTP_PASSWORD=${{ secrets.CI_GOTRUE_SMTP_PASS }}/' .env sed -i 's/APPFLOWY_AI_OPENAI_API_KEY=.*/APPFLOWY_AI_OPENAI_API_KEY=${{ secrets.CI_OPENAI_API_KEY }}/' .env - sed -i 's/APPFLOWY_OPENAI_API_KEY=.*/APPFLOWY_OPENAI_API_KEY=${{ secrets.CI_OPENAI_API_KEY }}/' .env - sed -i 's/APPFLOWY_INDEXER_OPENAI_API_KEY=.*/APPFLOWY_INDEXER_OPENAI_API_KEY=${{ secrets.CI_OPENAI_API_KEY }}/' .env sed -i 's/APPFLOWY_INDEXER_REDIS_URL=.*/APPFLOWY_INDEXER_REDIS_URL=redis:\/\/localhost:6379/' .env sed -i 's/APPFLOWY_INDEXER_DATABASE_URL=.*/APPFLOWY_INDEXER_DATABASE_URL=postgres:\/\/postgres:password@localhost:5432\/postgres/' .env diff --git a/Cargo.lock b/Cargo.lock index cf4eff1a..4e81bcd0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -637,7 +637,6 @@ dependencies = [ "log", "mime", "once_cell", - "openai_dive", "opener", "openssl", "pgvector", @@ -4017,22 +4016,6 @@ version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "624a8340c38c1b80fd549087862da4ba43e08858af025b236e509b6649fc13d5" -[[package]] -name = "openai_dive" -version = "0.4.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fbd2a1afd7e033794a61cd3d7b35474347309660979a163d5a77d7db67863b4d" -dependencies = [ - "bytes", - "reqwest 0.12.4", - "serde", - "serde_json", - "strum", - "strum_macros", - "tokio", - "tokio-util", -] - [[package]] name = "opener" version = "0.6.1" @@ -5076,7 +5059,6 @@ dependencies = [ "js-sys", "log", "mime", - "mime_guess", "native-tls", "once_cell", "percent-encoding", @@ -6207,25 +6189,6 @@ version = "0.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" -[[package]] -name = "strum" -version = "0.26.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5d8cec3501a5194c432b2b7976db6b7d10ec95c253208b45f83f7136aa985e29" - -[[package]] -name = "strum_macros" -version = "0.26.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c6cf59daf282c0a494ba14fd21610a0325f9f90ec9d1231dea26bcb1d696c946" -dependencies = [ - "heck 0.4.1", - "proc-macro2", - "quote", - "rustversion", - "syn 2.0.48", -] - [[package]] name = "subtle" version = "2.5.0" diff --git a/Cargo.toml b/Cargo.toml index 537892c4..8cd4b2d0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -73,7 +73,6 @@ appflowy-collaborate = { path = "services/appflowy-collaborate" } # ai appflowy-ai-client = { workspace = true, features = ["dto", "client-api"] } -openai_dive = { workspace = true, features = ["rustls-tls"] } pgvector = { workspace = true, features = ["sqlx"] } collab = { workspace = true } @@ -206,7 +205,6 @@ tonic = "0.11" prost = "0.12" tonic-proto = { path = "libs/tonic-proto" } appflowy-ai-client = { path = "libs/appflowy-ai-client", default-features = false } -openai_dive = { version = "0.4", features = ["rustls-tls"] } pgvector = { version = "0.3", features = ["sqlx"] } # collaboration diff --git a/docker-compose.yml b/docker-compose.yml index 40d8bf6a..65e72575 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -162,8 +162,8 @@ services: - APPFLOWY_INDEXER_REDIS_URL=redis://redis:6379 - APPFLOWY_INDEXER_ENVIRONMENT=production - APPFLOWY_INDEXER_DATABASE_URL=${APPFLOWY_INDEXER_DATABASE_URL} + - APPFLOWY_AI_SERVER_HOST=${APPFLOWY_AI_SERVER_HOST} - APPFLOWY_AI_SERVER_PORT=${APPFLOWY_AI_SERVER_PORT} - - APPFLOWY_AI_DATABASE_URL=${APPFLOWY_AI_DATABASE_URL} volumes: postgres_data: diff --git a/src/api/search.rs b/src/api/search.rs index 56768469..d2992af7 100644 --- a/src/api/search.rs +++ b/src/api/search.rs @@ -1,9 +1,7 @@ use actix_web::web::{Data, Query}; use actix_web::{web, Scope}; -use anyhow::anyhow; use uuid::Uuid; -use app_error::AppError; use authentication::jwt::Authorization; use shared_entity::dto::search_dto::{SearchDocumentRequest, SearchDocumentResponseItem}; use shared_entity::response::{AppResponse, JsonAppResponse}; @@ -26,18 +24,15 @@ async fn document_search( let request = payload.into_inner(); let user_uuid = auth.uuid()?; let uid = state.user_cache.get_user_uid(&user_uuid).await?; - let openai = match &state.openai { - Some(openai) => openai, - None => { - return Err( - AppError::Internal(anyhow!( - "Search API is not supported by this AppFlowy-Cloud instance" - )) - .into(), - ) - }, - }; let metrics = &*state.metrics.request_metrics; - let resp = search_document(&state.pg_pool, openai, uid, workspace_id, request, metrics).await?; + let resp = search_document( + &state.pg_pool, + &state.ai_client, + uid, + workspace_id, + request, + metrics, + ) + .await?; Ok(AppResponse::Ok().with_data(resp).into()) } diff --git a/src/application.rs b/src/application.rs index 28941c0e..f17b4a59 100644 --- a/src/application.rs +++ b/src/application.rs @@ -273,20 +273,9 @@ pub async fn init_state(config: &Config, rt_cmd_tx: CLCommandSender) -> Result Some(openai_dive::v1::api::Client::new( - key.expose_secret().clone(), - )), - _ => { - warn!("OpenAI API key not configured. Set APPFLOWY_OPENAI_API_KEY environment variable to enable OpenAI API."); - None - }, - }; - info!("Application state initialized"); Ok(AppState { pg_pool, - openai, config: Arc::new(config.clone()), user_cache, id_gen: Arc::new(RwLock::new(Snowflake::new(1))), diff --git a/src/biz/search/ops.rs b/src/biz/search/ops.rs index 96e024aa..b653241b 100644 --- a/src/biz/search/ops.rs +++ b/src/biz/search/ops.rs @@ -1,10 +1,10 @@ use crate::api::metrics::RequestMetrics; use app_error::ErrorCode; -use database::index::{search_documents, SearchDocumentParams}; -use openai_dive::v1::models::EmbeddingsEngine; -use openai_dive::v1::resources::embedding::{ - EmbeddingEncodingFormat, EmbeddingInput, EmbeddingOutput, EmbeddingParameters, +use appflowy_ai_client::client::AppFlowyAIClient; +use appflowy_ai_client::dto::{ + EmbeddingEncodingFormat, EmbeddingInput, EmbeddingOutput, EmbeddingRequest, EmbeddingsModel, }; +use database::index::{search_documents, SearchDocumentParams}; use shared_entity::dto::search_dto::{ SearchContentType, SearchDocumentRequest, SearchDocumentResponseItem, }; @@ -14,35 +14,29 @@ use uuid::Uuid; pub async fn search_document( pg_pool: &PgPool, - openai: &openai_dive::v1::api::Client, + openai: &AppFlowyAIClient, uid: i64, workspace_id: Uuid, request: SearchDocumentRequest, metrics: &RequestMetrics, ) -> Result, AppResponseError> { let embeddings = openai - .embeddings() - .create(EmbeddingParameters { + .embeddings(EmbeddingRequest { input: EmbeddingInput::String(request.query.clone()), - model: EmbeddingsEngine::TextEmbedding3Small.to_string(), - encoding_format: Some(EmbeddingEncodingFormat::Float), - dimensions: Some(1536), // text-embedding-3-small default number of dimensions - user: None, + model: EmbeddingsModel::TextEmbedding3Small.to_string(), + chunk_size: 0, + encoding_format: EmbeddingEncodingFormat::Float, + dimensions: 1536, }) .await .map_err(|e| AppResponseError::new(ErrorCode::Internal, e.to_string()))?; - - let tokens_used = if let Some(usage) = embeddings.usage { - metrics.record_search_tokens_used(&workspace_id, usage.total_tokens); - tracing::info!( - "workspace {} OpenAI API search tokens used: {}", - workspace_id, - usage.total_tokens - ); - usage.total_tokens - } else { - 0 - }; + let total_tokens = embeddings.total_tokens as u32; + metrics.record_search_tokens_used(&workspace_id, total_tokens); + tracing::info!( + "workspace {} OpenAI API search tokens used: {}", + workspace_id, + total_tokens + ); let embedding = embeddings .data @@ -71,7 +65,7 @@ pub async fn search_document( preview: request.preview_size.unwrap_or(180) as i32, embedding, }, - tokens_used, + total_tokens, ) .await?; tx.commit().await?; diff --git a/src/config/config.rs b/src/config/config.rs index 888be48e..ce61ea6f 100644 --- a/src/config/config.rs +++ b/src/config/config.rs @@ -14,7 +14,6 @@ pub struct Config { pub application: ApplicationSetting, pub websocket: WebsocketSetting, pub redis_uri: Secret, - pub openai_api_key: Option>, pub s3: S3Setting, pub appflowy_ai: AppFlowyAISetting, pub grpc_history: GrpcHistorySetting, @@ -161,9 +160,6 @@ pub fn get_configuration() -> Result { client_timeout: get_env_var("APPFLOWY_WEBSOCKET_CLIENT_TIMEOUT", "60").parse()?, }, redis_uri: get_env_var("APPFLOWY_REDIS_URI", "redis://localhost:6379").into(), - openai_api_key: std::env::var("APPFLOWY_OPENAI_API_KEY") - .map(Secret::from) - .ok(), s3: S3Setting { use_minio: get_env_var("APPFLOWY_S3_USE_MINIO", "true") .parse() diff --git a/src/state.rs b/src/state.rs index f084a562..e4049abd 100644 --- a/src/state.rs +++ b/src/state.rs @@ -48,7 +48,6 @@ pub struct AppState { pub metrics: AppMetrics, pub gotrue_admin: GoTrueAdmin, pub mailer: Mailer, - pub openai: Option, pub ai_client: AppFlowyAIClient, pub realtime_shared_state: RealtimeSharedState, pub grpc_history_client: Arc>>,