feat: reranking with fastembed added

html-router/src/html_state.rs

@@ -1,6 +1,7 @@
 use common::storage::db::SurrealDbClient;
 use common::utils::template_engine::{ProvidesTemplateEngine, TemplateEngine};
 use common::{create_template_engine, storage::db::ProvidesDb, utils::config::AppConfig};
+use composite_retrieval::reranking::RerankerPool;
 use std::sync::Arc;
 use tracing::debug;
 
@@ -13,6 +14,7 @@ pub struct HtmlState {
     pub templates: Arc<TemplateEngine>,
     pub session_store: Arc<SessionStoreType>,
     pub config: AppConfig,
+    pub reranker_pool: Option<Arc<RerankerPool>>,
 }
 
 impl HtmlState {
@@ -21,6 +23,7 @@ impl HtmlState {
         openai_client: Arc<OpenAIClientType>,
         session_store: Arc<SessionStoreType>,
         config: AppConfig,
+        reranker_pool: Option<Arc<RerankerPool>>,
     ) -> Result<Self, Box<dyn std::error::Error>> {
         let template_engine = create_template_engine!("templates");
         debug!("Template engine created for html_router.");
@@ -31,6 +34,7 @@ impl HtmlState {
             session_store,
             templates: Arc::new(template_engine),
             config,
+            reranker_pool,
         })
     }
 }

html-router/src/routes/chat/message_response_stream.rs

@@ -118,11 +118,17 @@ pub async fn get_response_stream(
         };
 
     // 2. Retrieve knowledge entities
+    let rerank_lease = match state.reranker_pool.as_ref() {
+        Some(pool) => Some(pool.checkout().await),
+        None => None,
+    };
+
     let entities = match retrieve_entities(
         &state.db,
         &state.openai_client,
         &user_message.content,
         &user.id,
+        rerank_lease,
     )
     .await
     {

html-router/src/routes/knowledge/handlers.rs

@@ -195,8 +195,19 @@ pub async fn suggest_knowledge_relationships(
 
     if !query_parts.is_empty() {
         let query = query_parts.join(" ");
-        if let Ok(results) =
-            retrieve_entities(&state.db, &state.openai_client, &query, &user.id).await
+        let rerank_lease = match state.reranker_pool.as_ref() {
+            Some(pool) => Some(pool.checkout().await),
+            None => None,
+        };
+
+        if let Ok(results) = retrieve_entities(
+            &state.db,
+            &state.openai_client,
+            &query,
+            &user.id,
+            rerank_lease,
+        )
+        .await
         {
             for RetrievedEntity { entity, score, .. } in results {
                 if suggestion_scores.len() >= MAX_RELATIONSHIP_SUGGESTIONS {