fix(index): return attr even when embedding failed

Signed-off-by: Wei Zhang <[email protected]>

crates/tabby-index/src/code/mod.rs

@@ -1,6 +1,6 @@
 use std::sync::Arc;
 
-use anyhow::{bail, Result};
+use anyhow::{anyhow, bail, Result};
 use async_stream::stream;
 use async_trait::async_trait;
 use futures::stream::BoxStream;
@@ -69,7 +69,7 @@ impl IndexAttributeBuilder<SourceCode> for CodeBuilder {
     async fn build_chunk_attributes<'a>(
         &self,
         source_code: &'a SourceCode,
-    ) -> BoxStream<'a, JoinHandle<Result<(Vec<String>, serde_json::Value)>>> {
+    ) -> BoxStream<'a, JoinHandle<(Result<Vec<String>>, serde_json::Value)>> {
         let text = match source_code.read_content() {
             Ok(content) => content,
             Err(e) => {
@@ -81,7 +81,7 @@ impl IndexAttributeBuilder<SourceCode> for CodeBuilder {
                 return Box::pin(stream! {
                     let path = source_code.filepath.clone();
                     yield tokio::spawn(async move {
-                        bail!("Failed to read content of '{}': {}", path, e);
+                        (Err(anyhow!("Failed to read content of '{}': {}", path, e)), json!({}))
                     });
                 });
             }
@@ -91,7 +91,7 @@ impl IndexAttributeBuilder<SourceCode> for CodeBuilder {
             warn!("No embedding service found for code indexing");
             return Box::pin(stream! {
                 yield tokio::spawn(async move {
-                    bail!("No embedding service found for code indexing");
+                    (Err(anyhow!("No embedding service found for code indexing")), json!({}))
                 });
             });
         };
@@ -110,10 +110,7 @@ impl IndexAttributeBuilder<SourceCode> for CodeBuilder {
                 let embedding = embedding.clone();
                 let rewritten_body = format!("```{}\n{}\n```", source_code.filepath, body);
                 yield tokio::spawn(async move {
-                    match build_binarize_embedding_tokens(embedding.clone(), &rewritten_body).await {
-                        Ok(tokens) => Ok((tokens, attributes)),
-                        Err(err) => Err(err),
-                    }
+                    (build_binarize_embedding_tokens(embedding.clone(), &rewritten_body).await, attributes)
                 });
             }
         };

crates/tabby-index/src/indexer.rs

@@ -43,7 +43,7 @@ pub trait IndexAttributeBuilder<T>: Send + Sync {
     async fn build_chunk_attributes<'a>(
         &self,
         document: &'a T,
-    ) -> BoxStream<'a, JoinHandle<Result<(Vec<String>, serde_json::Value)>>>;
+    ) -> BoxStream<'a, JoinHandle<(Result<Vec<String>>, serde_json::Value)>>;
 }
 
 pub struct TantivyDocBuilder<T> {
@@ -132,18 +132,11 @@ impl<T: ToIndexId> TantivyDocBuilder<T> {
                 // the document, and
                 // a flag indicating whether the tokens were created successfully.
                 yield tokio::spawn(async move {
-                    let Ok(built_chunk_attributes_result) = task.await else {
+                    let Ok((tokens, chunk_attributes)) = task.await else {
                         // Join error, there is no attr, return None and false
                         return (None, false);
                     };
 
-                    let (tokens, chunk_attributes) = match built_chunk_attributes_result{
-                        Ok((tokens, chunk_attributes)) => (tokens, chunk_attributes),
-                        Err(e) => {
-                            warn!("Failed to build chunk attributes for document '{}': {}", id, e);
-                            return (None, false);
-                        }
-                    };
                     let mut doc = doc! {
                         schema.field_id => id,
                         schema.field_source_id => source_id,
@@ -153,11 +146,18 @@ impl<T: ToIndexId> TantivyDocBuilder<T> {
                         schema.field_chunk_attributes => chunk_attributes,
                     };
 
-                    for token in &tokens {
-                        doc.add_text(schema.field_chunk_tokens, token);
+                    match tokens {
+                        Ok(tokens) => {
+                            for token in &tokens {
+                                doc.add_text(schema.field_chunk_tokens, token);
+                            }
+                            (Some(doc), true)
+                        },
+                        Err(e) => {
+                            warn!("Failed to build tokens for document '{}': {}", id, e);
+                            (Some(doc), false)
+                        }
                     }
-
-                    (Some(doc), !tokens.is_empty())
                 });
             }
         }

crates/tabby-index/src/structured_doc/mod.rs

@@ -38,7 +38,7 @@ impl IndexAttributeBuilder<StructuredDoc> for StructuredDocBuilder {
     async fn build_chunk_attributes<'a>(
         &self,
         document: &'a StructuredDoc,
-    ) -> BoxStream<'a, JoinHandle<Result<(Vec<String>, serde_json::Value)>>> {
+    ) -> BoxStream<'a, JoinHandle<(Result<Vec<String>>, serde_json::Value)>> {
         let embedding = self.embedding.clone();
         document.build_chunk_attributes(embedding).await
     }

crates/tabby-index/src/structured_doc/types.rs

@@ -53,7 +53,7 @@ pub trait BuildStructuredDoc {
     async fn build_chunk_attributes(
         &self,
         embedding: Arc<dyn Embedding>,
-    ) -> BoxStream<JoinHandle<Result<(Vec<String>, serde_json::Value)>>>;
+    ) -> BoxStream<JoinHandle<(Result<Vec<String>>, serde_json::Value)>>;
 }
 
 pub enum StructuredDocFields {
@@ -83,7 +83,7 @@ impl BuildStructuredDoc for StructuredDoc {
     async fn build_chunk_attributes(
         &self,
         embedding: Arc<dyn Embedding>,
-    ) -> BoxStream<JoinHandle<Result<(Vec<String>, serde_json::Value)>>> {
+    ) -> BoxStream<JoinHandle<(Result<Vec<String>>, serde_json::Value)>> {
         match &self.fields {
             StructuredDocFields::Web(doc) => doc.build_chunk_attributes(embedding).await,
             StructuredDocFields::Issue(doc) => doc.build_chunk_attributes(embedding).await,

crates/tabby-index/src/structured_doc/types/issue.rs

@@ -36,18 +36,13 @@ impl BuildStructuredDoc for IssueDocument {
     async fn build_chunk_attributes(
         &self,
         embedding: Arc<dyn Embedding>,
-    ) -> BoxStream<JoinHandle<Result<(Vec<String>, serde_json::Value)>>> {
+    ) -> BoxStream<JoinHandle<(Result<Vec<String>>, serde_json::Value)>> {
         let text = format!("{}\n\n{}", self.title, self.body);
         let s = stream! {
             yield tokio::spawn(async move {
-                let tokens = match build_tokens(embedding, &text).await{
-                    Ok(tokens) => tokens,
-                    Err(e) => {
-                        return Err(anyhow::anyhow!("Failed to build tokens for text: {}", e));
-                    }
-                };
+                let tokens = build_tokens(embedding, &text).await;
                 let chunk_attributes = json!({});
-                Ok((tokens, chunk_attributes))
+                (tokens, chunk_attributes)
             })
         };
 

crates/tabby-index/src/structured_doc/types/pull.rs

@@ -43,19 +43,13 @@ impl BuildStructuredDoc for PullDocument {
     async fn build_chunk_attributes(
         &self,
         embedding: Arc<dyn Embedding>,
-    ) -> BoxStream<JoinHandle<Result<(Vec<String>, serde_json::Value)>>> {
+    ) -> BoxStream<JoinHandle<(Result<Vec<String>>, serde_json::Value)>> {
         // currently not indexing the diff
         let text = format!("{}\n\n{}", self.title, self.body);
         let s = stream! {
             yield tokio::spawn(async move {
-                let tokens = match build_tokens(embedding, &text).await{
-                    Ok(tokens) => tokens,
-                    Err(e) => {
-                        return Err(anyhow::anyhow!("Failed to build tokens for text: {}", e));
-                    }
-                };
-                let chunk_attributes = json!({});
-                Ok((tokens, chunk_attributes))
+                let tokens = build_tokens(embedding, &text).await;
+                (tokens, json!({}))
             })
         };
 

crates/tabby-index/src/structured_doc/types/web.rs

@@ -34,7 +34,7 @@ impl BuildStructuredDoc for WebDocument {
     async fn build_chunk_attributes(
         &self,
         embedding: Arc<dyn Embedding>,
-    ) -> BoxStream<JoinHandle<Result<(Vec<String>, serde_json::Value)>>> {
+    ) -> BoxStream<JoinHandle<(Result<Vec<String>>, serde_json::Value)>> {
         let chunks: Vec<_> = TextSplitter::new(2048)
             .chunks(&self.body)
             .map(|x| x.to_owned())
@@ -45,7 +45,7 @@ impl BuildStructuredDoc for WebDocument {
             Err(e) => {
                 return Box::pin(stream! {
                     yield tokio::spawn(async move {
-                        Err(anyhow::anyhow!("Failed to build tokens for title: {}", e))
+                        (Err(anyhow::anyhow!("Failed to build tokens for title: {}", e)), json!({}))
                     });
                 });
             }
@@ -55,19 +55,19 @@ impl BuildStructuredDoc for WebDocument {
                 let title_embedding_tokens = title_embedding_tokens.clone();
                 let embedding = embedding.clone();
                 yield tokio::spawn(async move {
+                    let chunk = json!({
+                        fields::web::CHUNK_TEXT: chunk_text.clone(),
+                    });
                     let chunk_embedding_tokens = match build_tokens(embedding.clone(), &chunk_text).await {
                         Ok(tokens) => tokens,
                         Err(e) => {
-                            return Err(anyhow::anyhow!("Failed to build tokens for chunk: {}", e));
+                            return (Err(anyhow::anyhow!("Failed to build tokens for chunk: {}", e)), chunk);
                         }
                     };
-                    let chunk = json!({
-                        fields::web::CHUNK_TEXT: chunk_text,
-                    });
 
                     // Title embedding tokens are merged with chunk embedding tokens to enhance the search results.
                     let tokens = merge_tokens(vec![title_embedding_tokens, chunk_embedding_tokens]);
-                    Ok((tokens, chunk))
+                    (Ok(tokens), chunk)
                 });
             }
         };