From 7181bd5ffbf3193c0ebcda53e725fe00dd008fed Mon Sep 17 00:00:00 2001
From: Sebastian Hugentobler <shu@vanwa.ch>
Date: Tue, 1 Jul 2025 21:56:28 +0200
Subject: [PATCH] run text encoder functions in background thread for
 responsivness

---
 src/query.rs | 26 ++++++++++++++++++++++----
 1 file changed, 22 insertions(+), 4 deletions(-)

diff --git a/src/query.rs b/src/query.rs
index 599db3f..1d4e3a7 100644
--- a/src/query.rs
+++ b/src/query.rs
@@ -1,6 +1,7 @@
 //! Query processing and document retrieval.
 
 use snafu::{ResultExt, Snafu};
+use tokio::task::JoinError;
 
 use crate::{
     storage::{
@@ -18,12 +19,16 @@ pub enum AskError {
     Encode { source: tokenize::EncodeError },
     #[snafu(display("Failed to embed query."))]
     Embed { source: text_encoder::EmbedError },
+    #[snafu(display("Embedding task failed to execute."))]
+    EmbedJoin { source: JoinError },
     #[snafu(display("Failed to retrieve similar documents."))]
     Query {
         source: storage::queries::QueryError,
     },
     #[snafu(display("Failed to rerank documents."))]
     Rerank { source: text_encoder::RerankError },
+    #[snafu(display("Reranking task failed to execute."))]
+    RerankJoin { source: JoinError },
 }
 
 /// Process a user query and return ranked document matches.
@@ -37,15 +42,28 @@ pub async fn ask(
     limit: usize,
 ) -> Result<Vec<DocumentMatch>, AskError> {
     let encodings = tokenizer.encode(query, chunk_size).context(EncodeSnafu)?;
-    let embeddings = embedder.embed(encodings[0].clone()).context(EmbedSnafu)?;
+
+    let embedder = embedder.clone();
+    let encoding = encodings[0].clone();
+    let embeddings = tokio::task::spawn_blocking(move || embedder.embed(encoding))
+        .await
+        .context(EmbedJoinSnafu)?
+        .context(EmbedSnafu)?;
+
     let documents = db
         .query(embeddings, (limit * 10) as i32)
         .await
         .context(QuerySnafu)?;
 
-    let reranked_docs = reranker
-        .rerank(query, documents, tokenizer, limit)
-        .context(RerankSnafu)?;
+    let reranker = reranker.clone();
+    let tokenizer = tokenizer.clone();
+    let query = query.to_string();
+    let reranked_docs = tokio::task::spawn_blocking(move || {
+        reranker.rerank(&query, documents, &tokenizer, limit)
+    })
+    .await
+    .context(RerankJoinSnafu)?
+    .context(RerankSnafu)?;
 
     Ok(reranked_docs)
 }