feat: pool fastembed, batch embeddings, and reconcile embedding config on startup

ingestion-pipeline/src/pipeline/enrichment_result.rs

@@ -7,13 +7,12 @@ use serde::{Deserialize, Serialize};
 use common::{
     error::AppError,
     storage::{
-        db::SurrealDbClient,
         types::{
             knowledge_entity::{KnowledgeEntity, KnowledgeEntityType},
             knowledge_relationship::KnowledgeRelationship,
         },
     },
-    utils::{embedding::generate_embedding, embedding::EmbeddingProvider},
+    utils::embedding::EmbeddingProvider,
 };
 
 use crate::pipeline::context::EmbeddedKnowledgeEntity;
@@ -46,10 +45,8 @@ impl LLMEnrichmentResult {
         &self,
         source_id: &str,
         user_id: &str,
-        openai_client: &async_openai::Client<async_openai::config::OpenAIConfig>,
-        db_client: &SurrealDbClient,
         entity_concurrency: usize,
-        embedding_provider: Option<&EmbeddingProvider>,
+        embedding_provider: &EmbeddingProvider,
     ) -> Result<(Vec<EmbeddedKnowledgeEntity>, Vec<KnowledgeRelationship>), AppError> {
         let mapper = Arc::new(self.create_mapper());
 
@@ -58,8 +55,6 @@ impl LLMEnrichmentResult {
                 source_id,
                 user_id,
                 Arc::clone(&mapper),
-                openai_client,
-                db_client,
                 entity_concurrency,
                 embedding_provider,
             )
@@ -80,23 +75,18 @@ impl LLMEnrichmentResult {
         mapper
     }
 
-    #[allow(clippy::too_many_arguments)]
     async fn process_entities(
         &self,
         source_id: &str,
         user_id: &str,
         mapper: Arc<GraphMapper>,
-        openai_client: &async_openai::Client<async_openai::config::OpenAIConfig>,
-        db_client: &SurrealDbClient,
         entity_concurrency: usize,
-        embedding_provider: Option<&EmbeddingProvider>,
+        embedding_provider: &EmbeddingProvider,
     ) -> Result<Vec<EmbeddedKnowledgeEntity>, AppError> {
         stream::iter(self.knowledge_entities.clone().into_iter().map(|entity| {
             let mapper = Arc::clone(&mapper);
-            let openai_client = openai_client.clone();
             let source_id = source_id.to_string();
             let user_id = user_id.to_string();
-            let db_client = db_client.clone();
 
             async move {
                 create_single_entity(
@@ -104,8 +94,6 @@ impl LLMEnrichmentResult {
                     &source_id,
                     &user_id,
                     mapper,
-                    &openai_client,
-                    &db_client,
                     embedding_provider,
                 )
                 .await
@@ -145,9 +133,7 @@ async fn create_single_entity(
     source_id: &str,
     user_id: &str,
     mapper: Arc<GraphMapper>,
-    openai_client: &async_openai::Client<async_openai::config::OpenAIConfig>,
-    db_client: &SurrealDbClient,
-    embedding_provider: Option<&EmbeddingProvider>,
+    embedding_provider: &EmbeddingProvider,
 ) -> Result<EmbeddedKnowledgeEntity, AppError> {
     let assigned_id = mapper.get_id(&llm_entity.key)?.to_string();
 
@@ -156,11 +142,7 @@ async fn create_single_entity(
         llm_entity.name, llm_entity.description, llm_entity.entity_type
     );
 
-    let embedding = if let Some(provider) = embedding_provider {
-        provider.embed(&embedding_input).await?
-    } else {
-        generate_embedding(openai_client, &embedding_input, db_client).await?
-    };
+    let embedding = embedding_provider.embed(&embedding_input).await?;
 
     let now = Utc::now();
     let entity = KnowledgeEntity {

ingestion-pipeline/src/pipeline/services.rs

@@ -187,8 +187,7 @@ impl PipelineServices for DefaultPipelineServices {
         let config = retrieval_pipeline::RetrievalConfig::with_entities();
         match retrieval_pipeline::retrieve(
             &self.db,
-            &self.openai_client,
-            Some(&*self.embedding_provider),
+            &self.embedding_provider,
             &input_text,
             &content.user_id,
             config,
@@ -237,10 +236,8 @@ impl PipelineServices for DefaultPipelineServices {
             .to_database_entities(
                 content.id(),
                 &content.user_id,
-                &self.openai_client,
-                &self.db,
                 entity_concurrency,
-                Some(&*self.embedding_provider),
+                &self.embedding_provider,
             )
             .await
     }
@@ -258,15 +255,30 @@ impl PipelineServices for DefaultPipelineServices {
             overlap_tokens,
         )?;
 
+        if chunk_candidates.is_empty() {
+            return Ok(Vec::new());
+        }
+
+        // Embed all chunks of this document in one batch: a single lock acquisition and one
+        // blocking hop, letting the backend batch the inference internally.
+        let embeddings = self
+            .embedding_provider
+            .embed_batch(chunk_candidates.clone())
+            .await
+            .map_err(|e| {
+                AppError::InternalError(format!("FastEmbed embedding for chunks failed: {e}"))
+            })?;
+
+        if embeddings.len() != chunk_candidates.len() {
+            return Err(AppError::InternalError(format!(
+                "embedding batch returned {} vectors for {} chunks",
+                embeddings.len(),
+                chunk_candidates.len()
+            )));
+        }
+
         let mut chunks = Vec::with_capacity(chunk_candidates.len());
-        for chunk_text in chunk_candidates {
-            let embedding = self
-                .embedding_provider
-                .embed(&chunk_text)
-                .await
-                .map_err(|e| {
-                    AppError::InternalError(format!("FastEmbed embedding for chunk failed: {e}"))
-                })?;
+        for (chunk_text, embedding) in chunk_candidates.into_iter().zip(embeddings) {
             let chunk_struct = TextChunk::new(
                 content.id().to_string(),
                 chunk_text,