chore: use appflowy ai client for embeddings in search api

2024-06-18 10:51:11 +02:00 · 2024-06-18 10:51:11 +02:00 · 0e50650aa8
parent ae2a2a4fa0
commit 0e50650aa8
9 changed files with 28 additions and 96 deletions
--- 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

--- 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"
--- 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
--- 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:
--- 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())
 }
--- a/src/application.rs
+++ b/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))),
--- 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<Vec<SearchDocumentResponseItem>, 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?;
--- 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<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()
--- 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<openai_dive::v1::api::Client>,
  pub ai_client: AppFlowyAIClient,
  pub realtime_shared_state: RealtimeSharedState,
  pub grpc_history_client: Arc<Mutex<HistoryClient<tonic::transport::Channel>>>,