lunaris_retrieve/operators/
combinators.rs

1//! `.and / .or / .then` combinators.
2//!
3//! All three preserve the upstream operators as `Box<dyn Retriever>` so the
4//! downstream `fuse_rrf` operator (in `fuse.rs`) can introspect the tree
5//! shape AND so concurrent branches can run via `tokio::join!`.
6//!
7//! ## Semantics
8//!
9//! - `.and(A, B)` — runs A and B concurrently, **concatenates** their results
10//!   (per-source ranking preserved via `RawHit.source_op`). Downstream
11//!   `fuse_rrf` groups by `source_op` to fold the per-branch rankings.
12//! - `.or(A, B)` — runs A and B concurrently, **unions** their results by id;
13//!   on duplicate id, takes the max score. The merged hit's `source_op` is
14//!   the producer of the higher-scoring instance.
15//! - `.then(A, B)` — runs A first, then passes A's hit ids as a
16//!   `Filter::Or(Filter::Eq{id, …})` to B (B's results are the only ones
17//!   returned). Use this for "narrow then re-rank" pipelines.
18
19use std::any::Any;
20use std::collections::HashMap;
21
22use async_trait::async_trait;
23use lunaris_core::LunarisError;
24use lunaris_core::storage::types::Filter;
25
26use super::{QueryContext, Retriever};
27use crate::types::RawHit;
28
29/// `.and(A, B)` — concurrent fan-out, concatenated output (with per-source tag preserved).
30#[must_use = "AndRetriever is a query node — pass it to RetrievalBuilder::with_root() or wrap it via .fuse_rrf / .or, otherwise it never executes"]
31pub struct AndRetriever {
32    pub(crate) left: Box<dyn Retriever>,
33    pub(crate) right: Box<dyn Retriever>,
34}
35
36impl AndRetriever {
37    /// Read-only branch accessors for the documented `as_any` tree-inspection
38    /// pattern (operators/mod.rs) — lets external callers (e.g. harness shape
39    /// tests) walk a composed root without widening field visibility.
40    pub fn branches(&self) -> (&dyn Retriever, &dyn Retriever) {
41        (self.left.as_ref(), self.right.as_ref())
42    }
43}
44
45impl AndRetriever {
46    pub fn new(left: Box<dyn Retriever>, right: Box<dyn Retriever>) -> Self {
47        Self { left, right }
48    }
49
50    pub fn fuse_rrf(self, k: u32) -> super::fuse::FuseRrfRetriever {
51        super::fuse::FuseRrfRetriever::new(Box::new(self), k as usize)
52    }
53
54    pub fn top(self, n: usize) -> super::modifiers::TopRetriever {
55        super::modifiers::TopRetriever::new(Box::new(self), n)
56    }
57
58    /// Wrap with a cross-encoder rerank pass (Plan 02-03).
59    pub fn rerank(
60        self,
61        reranker: std::sync::Arc<dyn lunaris_rerank::Reranker>,
62    ) -> super::rerank::RerankRetriever {
63        super::rerank::RerankRetriever::new(Box::new(self), reranker)
64    }
65
66    /// Wrap with a fallback retriever — if THIS And path errors, switch to
67    /// `fallback` and tag returned hits with `degraded: true` (Plan 02-03).
68    pub fn degraded_fallback<R: Retriever + 'static>(
69        self,
70        fallback: R,
71    ) -> super::degraded::DegradedFallbackRetriever {
72        super::degraded::DegradedFallbackRetriever::new(Box::new(self), Box::new(fallback))
73    }
74}
75
76#[async_trait]
77impl Retriever for AndRetriever {
78    async fn retrieve(&self, ctx: &QueryContext) -> Result<Vec<RawHit>, LunarisError> {
79        let (left_res, right_res) = tokio::join!(self.left.retrieve(ctx), self.right.retrieve(ctx));
80        let mut out = left_res?;
81        out.extend(right_res?);
82        Ok(out)
83    }
84
85    fn as_any(&self) -> &dyn Any {
86        self
87    }
88}
89
90/// `.or(A, B)` — concurrent fan-out, unioned output (by id, max score).
91#[must_use = "OrRetriever is a query node — pass it to RetrievalBuilder::with_root() or wrap further, otherwise it never executes"]
92pub struct OrRetriever {
93    pub(crate) left: Box<dyn Retriever>,
94    pub(crate) right: Box<dyn Retriever>,
95}
96
97impl OrRetriever {
98    pub fn new(left: Box<dyn Retriever>, right: Box<dyn Retriever>) -> Self {
99        Self { left, right }
100    }
101
102    pub fn top(self, n: usize) -> super::modifiers::TopRetriever {
103        super::modifiers::TopRetriever::new(Box::new(self), n)
104    }
105}
106
107#[async_trait]
108impl Retriever for OrRetriever {
109    async fn retrieve(&self, ctx: &QueryContext) -> Result<Vec<RawHit>, LunarisError> {
110        let (left_res, right_res) = tokio::join!(self.left.retrieve(ctx), self.right.retrieve(ctx));
111        let mut by_id: HashMap<Vec<u8>, RawHit> = HashMap::new();
112        for h in left_res?.into_iter().chain(right_res?) {
113            by_id
114                .entry(h.id.clone())
115                .and_modify(|existing| {
116                    if h.score > existing.score {
117                        *existing = h.clone();
118                    }
119                })
120                .or_insert(h);
121        }
122        let mut out: Vec<RawHit> = by_id.into_values().collect();
123        // Stable order by descending score so .top(n) takes the top-N.
124        out.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(std::cmp::Ordering::Equal));
125        Ok(out)
126    }
127
128    fn as_any(&self) -> &dyn Any {
129        self
130    }
131}
132
133/// `.then(A, B)` — sequential narrow: A's ids constrain B via `Filter::Or(Eq{id})`.
134#[must_use = "ThenRetriever is a query node — pass it to RetrievalBuilder::with_root() or wrap further, otherwise it never executes"]
135pub struct ThenRetriever {
136    pub(crate) first: Box<dyn Retriever>,
137    pub(crate) second: Box<dyn Retriever>,
138}
139
140impl ThenRetriever {
141    pub fn new(first: Box<dyn Retriever>, second: Box<dyn Retriever>) -> Self {
142        Self { first, second }
143    }
144
145    pub fn top(self, n: usize) -> super::modifiers::TopRetriever {
146        super::modifiers::TopRetriever::new(Box::new(self), n)
147    }
148}
149
150/// Builder-friendly factory: `then(first, second)`. Mirrors the `rerank(upstream, reranker)`
151/// shape from `operators/rerank.rs` so callers can use function-style composition.
152pub fn then(first: Box<dyn Retriever>, second: Box<dyn Retriever>) -> ThenRetriever {
153    ThenRetriever::new(first, second)
154}
155
156#[async_trait]
157impl Retriever for ThenRetriever {
158    async fn retrieve(&self, ctx: &QueryContext) -> Result<Vec<RawHit>, LunarisError> {
159        let firsts = self.first.retrieve(ctx).await?;
160        if firsts.is_empty() {
161            return Ok(Vec::new());
162        }
163        // Build Filter::Or([Filter::Eq{field:"id", value:<id-as-string>}, ...])
164        let id_filter = Filter::Or(
165            firsts
166                .iter()
167                .map(|h| Filter::Eq {
168                    field: "id".to_string(),
169                    value: serde_json::Value::String(String::from_utf8_lossy(&h.id).into_owned()),
170                })
171                .collect(),
172        );
173
174        // Compose with any existing filter — Filter::And([existing, id_filter])
175        let new_filter = match &ctx.query.filter {
176            Some(existing) => Filter::And(vec![existing.clone(), id_filter]),
177            None => id_filter,
178        };
179
180        // Build a fresh ctx with the narrowed filter.
181        // Plan 260610-f91 (D-04): seed the narrowed OnceCell from the parent
182        // context if the parent already computed its embedding. This avoids a
183        // redundant embedder forward pass when both legs embed the same query
184        // text (e.g., then(Vector, Vector)). The query TEXT is unchanged by
185        // `then` — only the filter changes — so reuse is semantically exact.
186        // set() returns Err(value) only if the cell is already initialized,
187        // which is impossible here (brand-new cell). The discard is intentional.
188        let mut narrowed_query = ctx.query.clone();
189        narrowed_query.filter = Some(new_filter);
190
191        let narrowed_embedding = tokio::sync::OnceCell::new();
192        if let Some(existing) = ctx.query_embedding.get().cloned() {
193            let _ = narrowed_embedding.set(existing);
194        }
195
196        let narrowed_ctx = QueryContext {
197            query: narrowed_query,
198            scope: ctx.scope.clone(),
199            embedder: ctx.embedder.clone(),
200            storage: ctx.storage.clone(),
201            keyword: ctx.keyword.clone(),
202            query_embedding: narrowed_embedding,
203            moon_storage: ctx.moon_storage.clone(),
204        };
205
206        self.second.retrieve(&narrowed_ctx).await
207    }
208
209    fn as_any(&self) -> &dyn Any {
210        self
211    }
212}
lunaris_retrieve/operators/combinators.rs

lunaris_retrieve/operators/
combinators.rs
