retrieval-pipeline: v0

ingestion-pipeline/Cargo.toml

@@ -32,7 +32,7 @@ lopdf = "0.32"
 bytes = { workspace = true }
 
 common = { path = "../common" }
-composite-retrieval = { path = "../composite-retrieval" }
+retrieval-pipeline = { path = "../retrieval-pipeline" }
 async-trait = { workspace = true }
 state-machines = { workspace = true }
 [features]

ingestion-pipeline/src/pipeline/context.rs

@@ -11,7 +11,7 @@ use common::{
         },
     },
 };
-use composite_retrieval::RetrievedEntity;
+use retrieval_pipeline::RetrievedEntity;
 use tracing::error;
 
 use super::enrichment_result::LLMEnrichmentResult;

ingestion-pipeline/src/pipeline/mod.rs

@@ -28,7 +28,7 @@ use common::{
     },
     utils::config::AppConfig,
 };
-use composite_retrieval::reranking::RerankerPool;
+use retrieval_pipeline::reranking::RerankerPool;
 use tracing::{debug, info, warn};
 
 use self::{

ingestion-pipeline/src/pipeline/services.rs

@@ -19,8 +19,9 @@ use common::{
     },
     utils::{config::AppConfig, embedding::generate_embedding},
 };
-use composite_retrieval::{
-    reranking::RerankerPool, retrieve_entities, retrieved_entities_to_json, RetrievedEntity,
+use retrieval_pipeline::{
+    reranking::RerankerPool, retrieve_entities, retrieved_entities_to_json, RetrievalConfig,
+    RetrievalStrategy, RetrievedEntity, StrategyOutput,
 };
 use text_splitter::TextSplitter;
 
@@ -124,6 +125,14 @@ impl DefaultPipelineServices {
         Ok(request)
     }
 
+    fn configured_strategy(&self) -> RetrievalStrategy {
+        self.config
+            .retrieval_strategy
+            .as_deref()
+            .and_then(|value| value.parse().ok())
+            .unwrap_or(RetrievalStrategy::Initial)
+    }
+
     async fn perform_analysis(
         &self,
         request: CreateChatCompletionRequest,
@@ -178,14 +187,24 @@ impl PipelineServices for DefaultPipelineServices {
             None => None,
         };
 
-        retrieve_entities(
+        let mut config = RetrievalConfig::default();
+        config.strategy = self.configured_strategy();
+        match retrieve_entities(
             &self.db,
             &self.openai_client,
             &input_text,
             &content.user_id,
+            config,
             rerank_lease,
         )
         .await
+        {
+            Ok(StrategyOutput::Entities(entities)) => Ok(entities),
+            Ok(StrategyOutput::Chunks(_)) => Err(AppError::InternalError(
+                "Chunk-only retrieval is not supported in ingestion".into(),
+            )),
+            Err(err) => Err(err),
+        }
     }
 
     async fn run_enrichment(

ingestion-pipeline/src/pipeline/tests.rs

@@ -16,7 +16,7 @@ use common::{
         },
     },
 };
-use composite_retrieval::{RetrievedChunk, RetrievedEntity};
+use retrieval_pipeline::{RetrievedChunk, RetrievedEntity};
 use tokio::sync::Mutex;
 use uuid::Uuid;