chore: use appflowy ai client for embeddings in search api

.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
 

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"

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

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:

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())
 }

src/application.rs

@@ -273,20 +273,9 @@ pub async fn init_state(config: &Config, rt_cmd_tx: CLCommandSender) -> Result<A
     warn!("Failed to remove all connected users: {:?}", err);
   }
 
-  let openai = match &config.openai_api_key {
-    Some(key) if !key.expose_secret().is_empty() => 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))),

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 appflowy_ai_client::client::AppFlowyAIClient;
-use openai_dive::v1::models::EmbeddingsEngine;
+use appflowy_ai_client::dto::{
-use openai_dive::v1::resources::embedding::{
+  EmbeddingEncodingFormat, EmbeddingInput, EmbeddingOutput, EmbeddingRequest, EmbeddingsModel,
-  EmbeddingEncodingFormat, EmbeddingInput, EmbeddingOutput, EmbeddingParameters,
 };
+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<Vec<SearchDocumentResponseItem>, AppResponseError> {
   let embeddings = openai
-    .embeddings()
+    .embeddings(EmbeddingRequest {
-    .create(EmbeddingParameters {
       input: EmbeddingInput::String(request.query.clone()),
-      model: EmbeddingsEngine::TextEmbedding3Small.to_string(),
+      model: EmbeddingsModel::TextEmbedding3Small.to_string(),
-      encoding_format: Some(EmbeddingEncodingFormat::Float),
+      chunk_size: 0,
-      dimensions: Some(1536), // text-embedding-3-small default number of dimensions
+      encoding_format: EmbeddingEncodingFormat::Float,
-      user: None,
+      dimensions: 1536,
     })
     .await
     .map_err(|e| AppResponseError::new(ErrorCode::Internal, e.to_string()))?;
-
+  let total_tokens = embeddings.total_tokens as u32;
-  let tokens_used = if let Some(usage) = embeddings.usage {
+  metrics.record_search_tokens_used(&workspace_id, total_tokens);
-    metrics.record_search_tokens_used(&workspace_id, usage.total_tokens);
+  tracing::info!(
-    tracing::info!(
+    "workspace {} OpenAI API search tokens used: {}",
-      "workspace {} OpenAI API search tokens used: {}",
+    workspace_id,
-      workspace_id,
+    total_tokens
-      usage.total_tokens
+  );
-    );
-    usage.total_tokens
-  } else {
-    0
-  };
 
   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?;

src/config/config.rs

@@ -14,7 +14,6 @@ pub struct Config {
   pub application: ApplicationSetting,
   pub websocket: WebsocketSetting,
   pub redis_uri: Secret<String>,
-  pub openai_api_key: Option<Secret<String>>,
   pub s3: S3Setting,
   pub appflowy_ai: AppFlowyAISetting,
   pub grpc_history: GrpcHistorySetting,
@@ -161,9 +160,6 @@ pub fn get_configuration() -> Result<Config, anyhow::Error> {
       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()

src/state.rs

@@ -48,7 +48,6 @@ pub struct AppState {
   pub metrics: AppMetrics,
   pub gotrue_admin: GoTrueAdmin,
   pub mailer: Mailer,
-  pub openai: Option<openai_dive::v1::api::Client>,
   pub ai_client: AppFlowyAIClient,
   pub realtime_shared_state: RealtimeSharedState,
   pub grpc_history_client: Arc<Mutex<HistoryClient<tonic::transport::Channel>>>,