Skip to main content

lunaris_retrieve/
plan.rs

1//! F14 — build a real retriever tree from a JSON operator plan.
2//!
3//! The Python and TypeScript SDKs expose a composable DSL (`.and_()`,
4//! `.fuse_rrf()`, `.top()`) but their FFI historically carried a single flat
5//! `{index, k}` plan. Anything richer was collapsed to one leg, so both SDKs
6//! refused those plans rather than answer a different question than the one
7//! written. That refusal was right; this module removes the need for it.
8//!
9//! Both SDKs marshal their operator tree into the JSON shape below and call
10//! [`retriever_from_json`], so **one** parser decides what a plan means and
11//! the shape a caller writes is the shape the engine runs.
12//!
13//! ```text
14//! {"op":"vector",   "index":"chunks", "k":30}
15//! {"op":"keyword",  "index":"chunks", "k":30}
16//! {"op":"graph",    "seeds":[<seed>,…], "hops":2}
17//! {"op":"and",      "left":<node>, "right":<node>}
18//! {"op":"fuse_rrf", "k":60, "child":<node>}
19//! {"op":"top",      "n":5,  "child":<node>}
20//! ```
21//!
22//! A `<seed>` is either the 32-char lowercase hex an [`EntityId`] renders as
23//! (what the engine emits) or `{"name":"Alice","type":"Person"}` with an
24//! optional `"confidence"` (what a human writes). Both resolve through
25//! [`EntityId::from_name_and_type`] / [`EntityId::from_hex`] to the same
26//! anchor.
27//!
28//! ## Errors are the contract
29//!
30//! Every unrecognized op, missing branch and malformed seed is an
31//! [`PlanError`], never a skip and never a default. A parser that quietly
32//! drops a node it does not understand rebuilds the exact defect this module
33//! exists to remove: a plan that runs is not the plan that was written, and
34//! the caller gets a plausible list of hits with no indication of the swap.
35
36use lunaris_extract::types::EntityId;
37use serde_json::Value;
38
39use crate::operators::Retriever;
40use crate::operators::combinators::AndRetriever;
41use crate::operators::fuse::FuseRrfRetriever;
42use crate::operators::graph::Graph;
43use crate::operators::keyword::Keyword;
44use crate::operators::modifiers::TopRetriever;
45use crate::operators::vector::Vector;
46
47/// Why a JSON plan could not be turned into a retriever.
48#[derive(Debug, thiserror::Error)]
49pub enum PlanError {
50    #[error("plan node is not a JSON object: {0}")]
51    NotAnObject(String),
52    #[error("plan node has no `op` field: {0}")]
53    MissingOp(String),
54    #[error(
55        "unrecognized plan op `{0}` — the SDK plan parser does not build this operator, and \
56         skipping it would run a different plan than the one written"
57    )]
58    UnknownOp(String),
59    #[error("plan op `{op}` is missing required field `{field}`")]
60    MissingField { op: String, field: &'static str },
61    #[error("plan op `{op}` field `{field}` has the wrong type (wanted {wanted})")]
62    BadField { op: String, field: &'static str, wanted: &'static str },
63    #[error("graph seed {index} is neither 32-char hex nor a {{\"name\",\"type\"}} pair: {seed}")]
64    BadSeed { index: usize, seed: String },
65}
66
67type Built = Result<Box<dyn Retriever>, PlanError>;
68
69/// Build a retriever tree from a JSON plan node. See the module docs for the
70/// accepted shape.
71pub fn retriever_from_json(node: &Value) -> Built {
72    let obj = node.as_object().ok_or_else(|| PlanError::NotAnObject(node.to_string()))?;
73    let op = obj
74        .get("op")
75        .and_then(Value::as_str)
76        .ok_or_else(|| PlanError::MissingOp(node.to_string()))?;
77
78    match op {
79        "vector" => {
80            Ok(Box::new(Vector::new(str_field(node, op, "index")?, usize_field(node, op, "k")?)))
81        }
82        "keyword" => {
83            Ok(Box::new(Keyword::bm25(str_field(node, op, "index")?, usize_field(node, op, "k")?)))
84        }
85        "graph" => {
86            let seeds = seeds_field(node, op)?;
87            let hops = usize_field(node, op, "hops")?;
88            Ok(Box::new(Graph::anchored(seeds, hops)))
89        }
90        "and" => Ok(Box::new(AndRetriever::new(
91            retriever_from_json(child_field(node, op, "left")?)?,
92            retriever_from_json(child_field(node, op, "right")?)?,
93        ))),
94        "fuse_rrf" => Ok(Box::new(FuseRrfRetriever::new(
95            retriever_from_json(child_field(node, op, "child")?)?,
96            usize_field(node, op, "k")?,
97        ))),
98        "top" => Ok(Box::new(TopRetriever::new(
99            retriever_from_json(child_field(node, op, "child")?)?,
100            usize_field(node, op, "n")?,
101        ))),
102        other => Err(PlanError::UnknownOp(other.to_string())),
103    }
104}
105
106/// The hex `EntityId`s a graph root anchors on, or `None` when the root is
107/// not a [`Graph`]. Exists so a caller can prove the seeds it wrote are the
108/// seeds that were built — `plan_repr` deliberately renders only the seed
109/// COUNT, since the ids themselves are unbounded and belong in a trace, not
110/// in a plan string that gets compared for equality.
111pub fn seed_hex(r: &dyn Retriever) -> Option<Vec<String>> {
112    r.as_any()
113        .downcast_ref::<Graph>()
114        .map(|g| g.seeds.iter().map(|(id, _)| id.to_string()).collect())
115}
116
117fn field<'a>(node: &'a Value, op: &str, name: &'static str) -> Result<&'a Value, PlanError> {
118    node.get(name).ok_or_else(|| PlanError::MissingField { op: op.to_string(), field: name })
119}
120
121fn str_field<'a>(node: &'a Value, op: &str, name: &'static str) -> Result<&'a str, PlanError> {
122    field(node, op, name)?.as_str().ok_or_else(|| PlanError::BadField {
123        op: op.to_string(),
124        field: name,
125        wanted: "a string",
126    })
127}
128
129fn usize_field(node: &Value, op: &str, name: &'static str) -> Result<usize, PlanError> {
130    field(node, op, name)?.as_u64().map(|n| n as usize).ok_or_else(|| PlanError::BadField {
131        op: op.to_string(),
132        field: name,
133        wanted: "a non-negative integer",
134    })
135}
136
137fn child_field<'a>(node: &'a Value, op: &str, name: &'static str) -> Result<&'a Value, PlanError> {
138    field(node, op, name)
139}
140
141/// Parse `"seeds"` into the `(EntityId, confidence)` pairs `Graph::anchored`
142/// takes. A seed is a 32-char hex id or a `{"name","type"[,"confidence"]}`
143/// object; anything else is an error naming the offending index, because a
144/// dropped seed is an anchor the traversal silently never started from.
145fn seeds_field(node: &Value, op: &str) -> Result<Vec<(EntityId, f32)>, PlanError> {
146    let arr = field(node, op, "seeds")?.as_array().ok_or_else(|| PlanError::BadField {
147        op: op.to_string(),
148        field: "seeds",
149        wanted: "an array",
150    })?;
151    let mut out = Vec::with_capacity(arr.len());
152    for (index, seed) in arr.iter().enumerate() {
153        let bad = || PlanError::BadSeed { index, seed: seed.to_string() };
154        if let Some(s) = seed.as_str() {
155            out.push((EntityId::from_hex(s).ok_or_else(bad)?, 1.0));
156            continue;
157        }
158        let name = seed.get("name").and_then(Value::as_str).ok_or_else(bad)?;
159        let ty = seed.get("type").and_then(Value::as_str).ok_or_else(bad)?;
160        let conf = seed.get("confidence").and_then(Value::as_f64).unwrap_or(1.0) as f32;
161        out.push((EntityId::from_name_and_type(name, ty), conf));
162    }
163    Ok(out)
164}