retrieval-pipeline: v0

retrieval-pipeline/src/pipeline/config.rs

@@ -0,0 +1,121 @@
+use serde::{Deserialize, Serialize};
+use std::fmt;
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub enum RetrievalStrategy {
+    Initial,
+    Revised,
+}
+
+impl Default for RetrievalStrategy {
+    fn default() -> Self {
+        Self::Initial
+    }
+}
+
+impl std::str::FromStr for RetrievalStrategy {
+    type Err = String;
+
+    fn from_str(value: &str) -> Result<Self, Self::Err> {
+        match value.to_ascii_lowercase().as_str() {
+            "initial" => Ok(Self::Initial),
+            "revised" => Ok(Self::Revised),
+            other => Err(format!("unknown retrieval strategy '{other}'")),
+        }
+    }
+}
+
+impl fmt::Display for RetrievalStrategy {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let label = match self {
+            RetrievalStrategy::Initial => "initial",
+            RetrievalStrategy::Revised => "revised",
+        };
+        f.write_str(label)
+    }
+}
+
+/// Tunable parameters that govern each retrieval stage.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct RetrievalTuning {
+    pub entity_vector_take: usize,
+    pub chunk_vector_take: usize,
+    pub entity_fts_take: usize,
+    pub chunk_fts_take: usize,
+    pub score_threshold: f32,
+    pub fallback_min_results: usize,
+    pub token_budget_estimate: usize,
+    pub avg_chars_per_token: usize,
+    pub max_chunks_per_entity: usize,
+    pub lexical_match_weight: f32,
+    pub graph_traversal_seed_limit: usize,
+    pub graph_neighbor_limit: usize,
+    pub graph_score_decay: f32,
+    pub graph_seed_min_score: f32,
+    pub graph_vector_inheritance: f32,
+    pub rerank_blend_weight: f32,
+    pub rerank_scores_only: bool,
+    pub rerank_keep_top: usize,
+}
+
+impl Default for RetrievalTuning {
+    fn default() -> Self {
+        Self {
+            entity_vector_take: 15,
+            chunk_vector_take: 20,
+            entity_fts_take: 10,
+            chunk_fts_take: 20,
+            score_threshold: 0.35,
+            fallback_min_results: 10,
+            token_budget_estimate: 10000,
+            avg_chars_per_token: 4,
+            max_chunks_per_entity: 4,
+            lexical_match_weight: 0.15,
+            graph_traversal_seed_limit: 5,
+            graph_neighbor_limit: 6,
+            graph_score_decay: 0.75,
+            graph_seed_min_score: 0.4,
+            graph_vector_inheritance: 0.6,
+            rerank_blend_weight: 0.65,
+            rerank_scores_only: false,
+            rerank_keep_top: 8,
+        }
+    }
+}
+
+/// Wrapper containing tuning plus future flags for per-request overrides.
+#[derive(Debug, Clone)]
+pub struct RetrievalConfig {
+    pub strategy: RetrievalStrategy,
+    pub tuning: RetrievalTuning,
+}
+
+impl RetrievalConfig {
+    pub fn new(tuning: RetrievalTuning) -> Self {
+        Self {
+            strategy: RetrievalStrategy::Initial,
+            tuning,
+        }
+    }
+
+    pub fn with_strategy(strategy: RetrievalStrategy) -> Self {
+        Self {
+            strategy,
+            tuning: RetrievalTuning::default(),
+        }
+    }
+
+    pub fn with_tuning(strategy: RetrievalStrategy, tuning: RetrievalTuning) -> Self {
+        Self { strategy, tuning }
+    }
+}
+
+impl Default for RetrievalConfig {
+    fn default() -> Self {
+        Self {
+            strategy: RetrievalStrategy::default(),
+            tuning: RetrievalTuning::default(),
+        }
+    }
+}

retrieval-pipeline/src/pipeline/diagnostics.rs

@@ -0,0 +1,51 @@
+use serde::Serialize;
+
+/// Captures instrumentation for each hybrid retrieval stage when diagnostics are enabled.
+#[derive(Debug, Clone, Default, Serialize)]
+pub struct PipelineDiagnostics {
+    pub collect_candidates: Option<CollectCandidatesStats>,
+    pub enrich_chunks_from_entities: Option<ChunkEnrichmentStats>,
+    pub assemble: Option<AssembleStats>,
+}
+
+#[derive(Debug, Clone, Default, Serialize)]
+pub struct CollectCandidatesStats {
+    pub vector_entity_candidates: usize,
+    pub vector_chunk_candidates: usize,
+    pub fts_entity_candidates: usize,
+    pub fts_chunk_candidates: usize,
+    pub vector_chunk_scores: Vec<f32>,
+    pub fts_chunk_scores: Vec<f32>,
+}
+
+#[derive(Debug, Clone, Default, Serialize)]
+pub struct ChunkEnrichmentStats {
+    pub filtered_entity_count: usize,
+    pub fallback_min_results: usize,
+    pub chunk_sources_considered: usize,
+    pub chunk_candidates_before_enrichment: usize,
+    pub chunk_candidates_after_enrichment: usize,
+    pub top_chunk_scores: Vec<f32>,
+}
+
+#[derive(Debug, Clone, Default, Serialize)]
+pub struct AssembleStats {
+    pub token_budget_start: usize,
+    pub token_budget_spent: usize,
+    pub token_budget_remaining: usize,
+    pub budget_exhausted: bool,
+    pub chunks_selected: usize,
+    pub chunks_skipped_due_budget: usize,
+    pub entity_count: usize,
+    pub entity_traces: Vec<EntityAssemblyTrace>,
+}
+
+#[derive(Debug, Clone, Default, Serialize)]
+pub struct EntityAssemblyTrace {
+    pub entity_id: String,
+    pub source_id: String,
+    pub inspected_candidates: usize,
+    pub selected_chunk_ids: Vec<String>,
+    pub selected_chunk_scores: Vec<f32>,
+    pub skipped_due_budget: usize,
+}

retrieval-pipeline/src/pipeline/mod.rs

@@ -0,0 +1,397 @@
+mod config;
+mod diagnostics;
+mod stages;
+mod strategies;
+
+pub use config::{RetrievalConfig, RetrievalStrategy, RetrievalTuning};
+pub use diagnostics::{
+    AssembleStats, ChunkEnrichmentStats, CollectCandidatesStats, EntityAssemblyTrace,
+    PipelineDiagnostics,
+};
+
+use crate::{reranking::RerankerLease, RetrievedChunk, RetrievedEntity};
+use async_openai::Client;
+use async_trait::async_trait;
+use common::{error::AppError, storage::db::SurrealDbClient};
+use std::time::{Duration, Instant};
+use tracing::info;
+
+use stages::PipelineContext;
+use strategies::{InitialStrategyDriver, RevisedStrategyDriver};
+
+#[derive(Debug, Clone)]
+pub enum StrategyOutput {
+    Entities(Vec<RetrievedEntity>),
+    Chunks(Vec<RetrievedChunk>),
+}
+
+impl StrategyOutput {
+    pub fn as_entities(&self) -> Option<&[RetrievedEntity]> {
+        match self {
+            StrategyOutput::Entities(items) => Some(items),
+            _ => None,
+        }
+    }
+
+    pub fn into_entities(self) -> Option<Vec<RetrievedEntity>> {
+        match self {
+            StrategyOutput::Entities(items) => Some(items),
+            _ => None,
+        }
+    }
+
+    pub fn as_chunks(&self) -> Option<&[RetrievedChunk]> {
+        match self {
+            StrategyOutput::Chunks(items) => Some(items),
+            _ => None,
+        }
+    }
+
+    pub fn into_chunks(self) -> Option<Vec<RetrievedChunk>> {
+        match self {
+            StrategyOutput::Chunks(items) => Some(items),
+            _ => None,
+        }
+    }
+}
+
+#[derive(Debug)]
+pub struct PipelineRunOutput<T> {
+    pub results: T,
+    pub diagnostics: Option<PipelineDiagnostics>,
+    pub stage_timings: PipelineStageTimings,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum StageKind {
+    Embed,
+    CollectCandidates,
+    GraphExpansion,
+    ChunkAttach,
+    Rerank,
+    Assemble,
+}
+
+#[derive(Debug, Clone, Default, serde::Serialize)]
+pub struct PipelineStageTimings {
+    pub embed_ms: u128,
+    pub collect_candidates_ms: u128,
+    pub graph_expansion_ms: u128,
+    pub chunk_attach_ms: u128,
+    pub rerank_ms: u128,
+    pub assemble_ms: u128,
+}
+
+impl PipelineStageTimings {
+    pub fn record(&mut self, kind: StageKind, duration: Duration) {
+        let elapsed = duration.as_millis() as u128;
+        match kind {
+            StageKind::Embed => self.embed_ms += elapsed,
+            StageKind::CollectCandidates => self.collect_candidates_ms += elapsed,
+            StageKind::GraphExpansion => self.graph_expansion_ms += elapsed,
+            StageKind::ChunkAttach => self.chunk_attach_ms += elapsed,
+            StageKind::Rerank => self.rerank_ms += elapsed,
+            StageKind::Assemble => self.assemble_ms += elapsed,
+        }
+    }
+}
+
+#[async_trait]
+pub trait PipelineStage: Send + Sync {
+    fn kind(&self) -> StageKind;
+    async fn execute(&self, ctx: &mut PipelineContext<'_>) -> Result<(), AppError>;
+}
+
+pub type BoxedStage = Box<dyn PipelineStage + Send + Sync>;
+
+pub trait StrategyDriver {
+    type Output;
+
+    fn strategy(&self) -> RetrievalStrategy;
+    fn stages(&self) -> Vec<BoxedStage>;
+    fn override_tuning(&self, _config: &mut RetrievalConfig) {}
+
+    fn finalize(&self, ctx: &mut PipelineContext<'_>) -> Result<Self::Output, AppError>;
+}
+
+pub async fn run_pipeline(
+    db_client: &SurrealDbClient,
+    openai_client: &Client<async_openai::config::OpenAIConfig>,
+    input_text: &str,
+    user_id: &str,
+    config: RetrievalConfig,
+    reranker: Option<RerankerLease>,
+) -> Result<StrategyOutput, AppError> {
+    let input_chars = input_text.chars().count();
+    let input_preview: String = input_text.chars().take(120).collect();
+    let input_preview_clean = input_preview.replace('\n', " ");
+    let preview_len = input_preview_clean.chars().count();
+    info!(
+        %user_id,
+        input_chars,
+        preview_truncated = input_chars > preview_len,
+        preview = %input_preview_clean,
+        "Starting ingestion retrieval pipeline"
+    );
+
+    if config.strategy == RetrievalStrategy::Initial {
+        let driver = InitialStrategyDriver::new();
+        let run = execute_strategy(
+            driver,
+            db_client,
+            openai_client,
+            None,
+            input_text,
+            user_id,
+            config,
+            reranker,
+            false,
+        )
+        .await?;
+        return Ok(StrategyOutput::Entities(run.results));
+    }
+
+    let driver = RevisedStrategyDriver::new();
+    let run = execute_strategy(
+        driver,
+        db_client,
+        openai_client,
+        None,
+        input_text,
+        user_id,
+        config,
+        reranker,
+        false,
+    )
+    .await?;
+    Ok(StrategyOutput::Chunks(run.results))
+}
+
+pub async fn run_pipeline_with_embedding(
+    db_client: &SurrealDbClient,
+    openai_client: &Client<async_openai::config::OpenAIConfig>,
+    query_embedding: Vec<f32>,
+    input_text: &str,
+    user_id: &str,
+    config: RetrievalConfig,
+    reranker: Option<RerankerLease>,
+) -> Result<StrategyOutput, AppError> {
+    if config.strategy == RetrievalStrategy::Initial {
+        let driver = InitialStrategyDriver::new();
+        let run = execute_strategy(
+            driver,
+            db_client,
+            openai_client,
+            Some(query_embedding),
+            input_text,
+            user_id,
+            config,
+            reranker,
+            false,
+        )
+        .await?;
+        return Ok(StrategyOutput::Entities(run.results));
+    }
+
+    let driver = RevisedStrategyDriver::new();
+    let run = execute_strategy(
+        driver,
+        db_client,
+        openai_client,
+        Some(query_embedding),
+        input_text,
+        user_id,
+        config,
+        reranker,
+        false,
+    )
+    .await?;
+    Ok(StrategyOutput::Chunks(run.results))
+}
+
+pub async fn run_pipeline_with_embedding_with_metrics(
+    db_client: &SurrealDbClient,
+    openai_client: &Client<async_openai::config::OpenAIConfig>,
+    query_embedding: Vec<f32>,
+    input_text: &str,
+    user_id: &str,
+    config: RetrievalConfig,
+    reranker: Option<RerankerLease>,
+) -> Result<PipelineRunOutput<StrategyOutput>, AppError> {
+    if config.strategy == RetrievalStrategy::Initial {
+        let driver = InitialStrategyDriver::new();
+        let run = execute_strategy(
+            driver,
+            db_client,
+            openai_client,
+            Some(query_embedding),
+            input_text,
+            user_id,
+            config,
+            reranker,
+            false,
+        )
+        .await?;
+        return Ok(PipelineRunOutput {
+            results: StrategyOutput::Entities(run.results),
+            diagnostics: run.diagnostics,
+            stage_timings: run.stage_timings,
+        });
+    }
+
+    let driver = RevisedStrategyDriver::new();
+    let run = execute_strategy(
+        driver,
+        db_client,
+        openai_client,
+        Some(query_embedding),
+        input_text,
+        user_id,
+        config,
+        reranker,
+        false,
+    )
+    .await?;
+    Ok(PipelineRunOutput {
+        results: StrategyOutput::Chunks(run.results),
+        diagnostics: run.diagnostics,
+        stage_timings: run.stage_timings,
+    })
+}
+
+pub async fn run_pipeline_with_embedding_with_diagnostics(
+    db_client: &SurrealDbClient,
+    openai_client: &Client<async_openai::config::OpenAIConfig>,
+    query_embedding: Vec<f32>,
+    input_text: &str,
+    user_id: &str,
+    config: RetrievalConfig,
+    reranker: Option<RerankerLease>,
+) -> Result<PipelineRunOutput<StrategyOutput>, AppError> {
+    if config.strategy == RetrievalStrategy::Initial {
+        let driver = InitialStrategyDriver::new();
+        let run = execute_strategy(
+            driver,
+            db_client,
+            openai_client,
+            Some(query_embedding),
+            input_text,
+            user_id,
+            config,
+            reranker,
+            true,
+        )
+        .await?;
+        return Ok(PipelineRunOutput {
+            results: StrategyOutput::Entities(run.results),
+            diagnostics: run.diagnostics,
+            stage_timings: run.stage_timings,
+        });
+    }
+
+    let driver = RevisedStrategyDriver::new();
+    let run = execute_strategy(
+        driver,
+        db_client,
+        openai_client,
+        Some(query_embedding),
+        input_text,
+        user_id,
+        config,
+        reranker,
+        true,
+    )
+    .await?;
+    Ok(PipelineRunOutput {
+        results: StrategyOutput::Chunks(run.results),
+        diagnostics: run.diagnostics,
+        stage_timings: run.stage_timings,
+    })
+}
+
+pub fn retrieved_entities_to_json(entities: &[RetrievedEntity]) -> serde_json::Value {
+    serde_json::json!(entities
+        .iter()
+        .map(|entry| {
+            serde_json::json!({
+                "KnowledgeEntity": {
+                    "id": entry.entity.id,
+                    "name": entry.entity.name,
+                    "description": entry.entity.description,
+                    "score": round_score(entry.score),
+                    "chunks": entry.chunks.iter().map(|chunk| {
+                        serde_json::json!({
+                            "score": round_score(chunk.score),
+                            "content": chunk.chunk.chunk
+                        })
+                    }).collect::<Vec<_>>()
+                }
+            })
+        })
+        .collect::<Vec<_>>())
+}
+
+async fn execute_strategy<D: StrategyDriver>(
+    driver: D,
+    db_client: &SurrealDbClient,
+    openai_client: &Client<async_openai::config::OpenAIConfig>,
+    query_embedding: Option<Vec<f32>>,
+    input_text: &str,
+    user_id: &str,
+    mut config: RetrievalConfig,
+    reranker: Option<RerankerLease>,
+    capture_diagnostics: bool,
+) -> Result<PipelineRunOutput<D::Output>, AppError> {
+    driver.override_tuning(&mut config);
+    let ctx = match query_embedding {
+        Some(embedding) => PipelineContext::with_embedding(
+            db_client,
+            openai_client,
+            embedding,
+            input_text.to_owned(),
+            user_id.to_owned(),
+            config,
+            reranker,
+        ),
+        None => PipelineContext::new(
+            db_client,
+            openai_client,
+            input_text.to_owned(),
+            user_id.to_owned(),
+            config,
+            reranker,
+        ),
+    };
+
+    run_with_driver(driver, ctx, capture_diagnostics).await
+}
+
+async fn run_with_driver<D: StrategyDriver>(
+    driver: D,
+    mut ctx: PipelineContext<'_>,
+    capture_diagnostics: bool,
+) -> Result<PipelineRunOutput<D::Output>, AppError> {
+    if capture_diagnostics {
+        ctx.enable_diagnostics();
+    }
+
+    for stage in driver.stages() {
+        let start = Instant::now();
+        stage.execute(&mut ctx).await?;
+        ctx.record_stage_duration(stage.kind(), start.elapsed());
+    }
+
+    let diagnostics = ctx.take_diagnostics();
+    let stage_timings = ctx.take_stage_timings();
+    let results = driver.finalize(&mut ctx)?;
+
+    Ok(PipelineRunOutput {
+        results,
+        diagnostics,
+        stage_timings,
+    })
+}
+
+fn round_score(value: f32) -> f64 {
+    (f64::from(value) * 1000.0).round() / 1000.0
+}

retrieval-pipeline/src/pipeline/strategies.rs

@@ -0,0 +1,75 @@
+use super::{
+    stages::{
+        AssembleEntitiesStage, ChunkAssembleStage, ChunkAttachStage, ChunkRerankStage,
+        ChunkVectorStage, CollectCandidatesStage, EmbedStage, GraphExpansionStage, PipelineContext,
+        RerankStage,
+    },
+    BoxedStage, RetrievalConfig, RetrievalStrategy, StrategyDriver,
+};
+use crate::{RetrievedChunk, RetrievedEntity};
+use common::error::AppError;
+
+pub struct InitialStrategyDriver;
+
+impl InitialStrategyDriver {
+    pub fn new() -> Self {
+        Self
+    }
+}
+
+impl StrategyDriver for InitialStrategyDriver {
+    type Output = Vec<RetrievedEntity>;
+
+    fn strategy(&self) -> RetrievalStrategy {
+        RetrievalStrategy::Initial
+    }
+
+    fn stages(&self) -> Vec<BoxedStage> {
+        vec![
+            Box::new(EmbedStage),
+            Box::new(CollectCandidatesStage),
+            Box::new(GraphExpansionStage),
+            Box::new(ChunkAttachStage),
+            Box::new(RerankStage),
+            Box::new(AssembleEntitiesStage),
+        ]
+    }
+
+    fn finalize(&self, ctx: &mut PipelineContext<'_>) -> Result<Self::Output, AppError> {
+        Ok(ctx.take_entity_results())
+    }
+}
+
+pub struct RevisedStrategyDriver;
+
+impl RevisedStrategyDriver {
+    pub fn new() -> Self {
+        Self
+    }
+}
+
+impl StrategyDriver for RevisedStrategyDriver {
+    type Output = Vec<RetrievedChunk>;
+
+    fn strategy(&self) -> RetrievalStrategy {
+        RetrievalStrategy::Revised
+    }
+
+    fn stages(&self) -> Vec<BoxedStage> {
+        vec![
+            Box::new(EmbedStage),
+            Box::new(ChunkVectorStage),
+            Box::new(ChunkRerankStage),
+            Box::new(ChunkAssembleStage),
+        ]
+    }
+
+    fn finalize(&self, ctx: &mut PipelineContext<'_>) -> Result<Self::Output, AppError> {
+        Ok(ctx.take_chunk_results())
+    }
+
+    fn override_tuning(&self, config: &mut RetrievalConfig) {
+        config.tuning.entity_vector_take = 0;
+        config.tuning.entity_fts_take = 0;
+    }
+}