retrieval-pipeline: v0

html-router/Cargo.toml

@@ -38,7 +38,7 @@ url = { workspace = true }
 uuid = { workspace = true }
 
 common = { path = "../common" }
-composite-retrieval = { path = "../composite-retrieval" }
+retrieval-pipeline = { path = "../retrieval-pipeline" }
 json-stream-parser = { path = "../json-stream-parser" }
 
 [build-dependencies]

html-router/src/html_state.rs

@@ -1,7 +1,7 @@
 use common::storage::{db::SurrealDbClient, store::StorageManager};
 use common::utils::template_engine::{ProvidesTemplateEngine, TemplateEngine};
 use common::{create_template_engine, storage::db::ProvidesDb, utils::config::AppConfig};
-use composite_retrieval::reranking::RerankerPool;
+use retrieval_pipeline::{reranking::RerankerPool, RetrievalStrategy};
 use std::sync::Arc;
 use tracing::debug;
 
@@ -40,6 +40,14 @@ impl HtmlState {
             reranker_pool,
         })
     }
+
+    pub fn retrieval_strategy(&self) -> RetrievalStrategy {
+        self.config
+            .retrieval_strategy
+            .as_deref()
+            .and_then(|value| value.parse().ok())
+            .unwrap_or(RetrievalStrategy::Initial)
+    }
 }
 impl ProvidesDb for HtmlState {
     fn db(&self) -> &Arc<SurrealDbClient> {

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

@@ -8,16 +8,16 @@ use axum::{
         Sse,
     },
 };
-use composite_retrieval::{
-    answer_retrieval::{create_chat_request, create_user_message_with_history, LLMResponseFormat},
-    retrieve_entities, retrieved_entities_to_json,
-};
 use futures::{
     stream::{self, once},
     Stream, StreamExt, TryStreamExt,
 };
 use json_stream_parser::JsonStreamParser;
 use minijinja::Value;
+use retrieval_pipeline::{
+    answer_retrieval::{create_chat_request, create_user_message_with_history, LLMResponseFormat},
+    retrieve_entities, retrieved_entities_to_json, RetrievalConfig, StrategyOutput,
+};
 use serde::{Deserialize, Serialize};
 use serde_json::from_str;
 use tokio::sync::{mpsc::channel, Mutex};
@@ -123,16 +123,24 @@ pub async fn get_response_stream(
         None => None,
     };
 
+    let mut retrieval_config = RetrievalConfig::default();
+    retrieval_config.strategy = state.retrieval_strategy();
     let entities = match retrieve_entities(
         &state.db,
         &state.openai_client,
         &user_message.content,
         &user.id,
+        retrieval_config,
         rerank_lease,
     )
     .await
     {
-        Ok(entities) => entities,
+        Ok(StrategyOutput::Entities(entities)) => entities,
+        Ok(StrategyOutput::Chunks(_)) => {
+            return Sse::new(create_error_stream(
+                "Chunk-only retrieval results are not supported in this route",
+            ))
+        }
         Err(_e) => {
             return Sse::new(create_error_stream("Failed to retrieve knowledge entities"));
         }

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

@@ -24,7 +24,7 @@ use common::{
     },
     utils::embedding::generate_embedding,
 };
-use composite_retrieval::{retrieve_entities, RetrievedEntity};
+use retrieval_pipeline::{retrieve_entities, RetrievalConfig, RetrievedEntity, StrategyOutput};
 use tracing::debug;
 use uuid::Uuid;
 
@@ -284,11 +284,15 @@ pub async fn suggest_knowledge_relationships(
             None => None,
         };
 
-        if let Ok(results) = retrieve_entities(
+        let mut retrieval_config = RetrievalConfig::default();
+        retrieval_config.strategy = state.retrieval_strategy();
+
+        if let Ok(StrategyOutput::Entities(results)) = retrieve_entities(
             &state.db,
             &state.openai_client,
             &query,
             &user.id,
+            retrieval_config,
             rerank_lease,
         )
         .await