chore: adjust redis env (#1003)

AppFlowy-IO · Nov 17, 2024 · fd98351 · fd98351
1 parent 3799966
commit fd98351
Show file tree

Hide file tree

Showing 7 changed files with 14 additions and 5 deletions.
diff --git a/.github/workflows/integration_test.yml b/.github/workflows/integration_test.yml
@@ -24,6 +24,7 @@ env:
   LOCALHOST_URL: http://localhost
   LOCALHOST_WS: ws://localhost/ws/v1
   APPFLOWY_REDIS_URI: redis://redis:6379
+  APPFLOWY_AI_REDIS_URL: redis://redis:6379
   LOCALHOST_GOTRUE: http://localhost/gotrue
   POSTGRES_PASSWORD: password
   DATABASE_URL: postgres://postgres:password@localhost:5432/postgres

diff --git a/deploy.env b/deploy.env
@@ -146,6 +146,7 @@ APPFLOWY_AI_OPENAI_API_KEY=
 APPFLOWY_AI_SERVER_PORT=5001
 APPFLOWY_AI_SERVER_HOST=ai
 APPFLOWY_AI_DATABASE_URL=postgresql+psycopg://${POSTGRES_USER}:${POSTGRES_PASSWORD}@${POSTGRES_HOST}:${POSTGRES_PORT}/${POSTGRES_DB}
+APPFLOWY_AI_REDIS_URL=redis://${REDIS_HOST}:${REDIS_PORT}
 APPFLOWY_LOCAL_AI_TEST_ENABLED=false
 
 # AppFlowy Indexer

diff --git a/dev.env b/dev.env
@@ -113,6 +113,7 @@ APPFLOWY_AI_OPENAI_API_KEY=
 APPFLOWY_AI_SERVER_PORT=5001
 APPFLOWY_AI_SERVER_HOST=localhost
 APPFLOWY_AI_DATABASE_URL=postgresql+psycopg://postgres:password@postgres:5432/postgres
+APPFLOWY_AI_REDIS_URL=redis://redis:6379
 APPFLOWY_LOCAL_AI_TEST_ENABLED=false
 
 # AppFlowy Indexer

diff --git a/docker-compose-ci.yml b/docker-compose-ci.yml
@@ -148,7 +148,7 @@ services:
       - LOCAL_AI_AWS_SECRET_ACCESS_KEY=${LOCAL_AI_AWS_SECRET_ACCESS_KEY}
       - APPFLOWY_AI_SERVER_PORT=${APPFLOWY_AI_SERVER_PORT}
       - APPFLOWY_AI_DATABASE_URL=${APPFLOWY_AI_DATABASE_URL}
-      - APPFLOWY_AI_REDIS_URL=${APPFLOWY_REDIS_URI}
+      - APPFLOWY_AI_REDIS_URL=${APPFLOWY_AI_REDIS_URL}
 
   appflowy_worker:
     restart: on-failure

diff --git a/docker-compose-dev.yml b/docker-compose-dev.yml
@@ -122,6 +122,7 @@ services:
       - OPENAI_API_KEY=${APPFLOWY_AI_OPENAI_API_KEY}
       - APPFLOWY_AI_SERVER_PORT=${APPFLOWY_AI_SERVER_PORT}
       - APPFLOWY_AI_DATABASE_URL=${APPFLOWY_AI_DATABASE_URL}
+      - APPFLOWY_AI_REDIS_URL=${APPFLOWY_AI_REDIS_URL}
 
 volumes:
   postgres_data:
diff --git a/docker-compose.yml b/docker-compose.yml
@@ -140,6 +140,7 @@ services:
       - OPENAI_API_KEY=${APPFLOWY_AI_OPENAI_API_KEY}
       - APPFLOWY_AI_SERVER_PORT=${APPFLOWY_AI_SERVER_PORT}
       - APPFLOWY_AI_DATABASE_URL=${APPFLOWY_AI_DATABASE_URL}
+      - APPFLOWY_AI_REDIS_URL=${APPFLOWY_AI_REDIS_URL}
 
   appflowy_worker:
     restart: on-failure

diff --git a/services/appflowy-collaborate/src/indexer/document_indexer.rs b/services/appflowy-collaborate/src/indexer/document_indexer.rs
@@ -25,15 +25,21 @@ pub struct DocumentIndexer {
   ai_client: AppFlowyAIClient,
   tokenizer: Arc<CoreBPE>,
   embedding_model: EmbeddingModel,
+  use_tiktoken: bool,
 }
 
 impl DocumentIndexer {
   pub fn new(ai_client: AppFlowyAIClient) -> Arc<Self> {
     let tokenizer = tiktoken_rs::cl100k_base().unwrap();
+    let use_tiktoken = get_env_var("APPFLOWY_AI_CONTENT_SPLITTER_TIKTOKEN", "false")
+      .parse::<bool>()
+      .unwrap_or(false);
+
     Arc::new(Self {
       ai_client,
       tokenizer: Arc::new(tokenizer),
       embedding_model: EmbeddingModel::TextEmbedding3Small,
+      use_tiktoken,
     })
   }
 }
@@ -62,6 +68,7 @@ impl Indexer for DocumentIndexer {
           CollabType::Document,
           &self.embedding_model,
           self.tokenizer.clone(),
+          self.use_tiktoken,
         )
         .await
       },
@@ -136,11 +143,8 @@ async fn create_embedding(
   collab_type: CollabType,
   embedding_model: &EmbeddingModel,
   tokenizer: Arc<CoreBPE>,
+  use_tiktoken: bool,
 ) -> Result<Vec<AFCollabEmbeddingParams>, AppError> {
-  let use_tiktoken = get_env_var("APPFLOWY_AI_CONTENT_SPLITTER_TIKTOKEN", "false")
-    .parse::<bool>()
-    .unwrap_or(false);
-
   let split_contents = if use_tiktoken {
     let max_tokens = embedding_model.default_dimensions() as usize;
     if content.len() < 500 {