use lunaris_extract::types::EntityId;
use serde_json::Value;
use crate::operators::Retriever;
use crate::operators::combinators::AndRetriever;
use crate::operators::fuse::FuseRrfRetriever;
use crate::operators::graph::Graph;
use crate::operators::keyword::Keyword;
use crate::operators::modifiers::TopRetriever;
use crate::operators::vector::Vector;
#[derive(Debug, thiserror::Error)]
pub enum PlanError {
#[error("plan node is not a JSON object: {0}")]
NotAnObject(String),
#[error("plan node has no `op` field: {0}")]
MissingOp(String),
#[error(
"unrecognized plan op `{0}` — the SDK plan parser does not build this operator, and \
skipping it would run a different plan than the one written"
)]
UnknownOp(String),
#[error("plan op `{op}` is missing required field `{field}`")]
MissingField { op: String, field: &'static str },
#[error("plan op `{op}` field `{field}` has the wrong type (wanted {wanted})")]
BadField { op: String, field: &'static str, wanted: &'static str },
#[error("graph seed {index} is neither 32-char hex nor a {{\"name\",\"type\"}} pair: {seed}")]
BadSeed { index: usize, seed: String },
}
type Built = Result<Box<dyn Retriever>, PlanError>;
pub fn retriever_from_json(node: &Value) -> Built {
let obj = node.as_object().ok_or_else(|| PlanError::NotAnObject(node.to_string()))?;
let op = obj
.get("op")
.and_then(Value::as_str)
.ok_or_else(|| PlanError::MissingOp(node.to_string()))?;
match op {
"vector" => {
Ok(Box::new(Vector::new(str_field(node, op, "index")?, usize_field(node, op, "k")?)))
}
"keyword" => {
Ok(Box::new(Keyword::bm25(str_field(node, op, "index")?, usize_field(node, op, "k")?)))
}
"graph" => {
let seeds = seeds_field(node, op)?;
let hops = usize_field(node, op, "hops")?;
Ok(Box::new(Graph::anchored(seeds, hops)))
}
"and" => Ok(Box::new(AndRetriever::new(
retriever_from_json(child_field(node, op, "left")?)?,
retriever_from_json(child_field(node, op, "right")?)?,
))),
"fuse_rrf" => Ok(Box::new(FuseRrfRetriever::new(
retriever_from_json(child_field(node, op, "child")?)?,
usize_field(node, op, "k")?,
))),
"top" => Ok(Box::new(TopRetriever::new(
retriever_from_json(child_field(node, op, "child")?)?,
usize_field(node, op, "n")?,
))),
other => Err(PlanError::UnknownOp(other.to_string())),
}
}
pub fn seed_hex(r: &dyn Retriever) -> Option<Vec<String>> {
r.as_any()
.downcast_ref::<Graph>()
.map(|g| g.seeds.iter().map(|(id, _)| id.to_string()).collect())
}
fn field<'a>(node: &'a Value, op: &str, name: &'static str) -> Result<&'a Value, PlanError> {
node.get(name).ok_or_else(|| PlanError::MissingField { op: op.to_string(), field: name })
}
fn str_field<'a>(node: &'a Value, op: &str, name: &'static str) -> Result<&'a str, PlanError> {
field(node, op, name)?.as_str().ok_or_else(|| PlanError::BadField {
op: op.to_string(),
field: name,
wanted: "a string",
})
}
fn usize_field(node: &Value, op: &str, name: &'static str) -> Result<usize, PlanError> {
field(node, op, name)?.as_u64().map(|n| n as usize).ok_or_else(|| PlanError::BadField {
op: op.to_string(),
field: name,
wanted: "a non-negative integer",
})
}
fn child_field<'a>(node: &'a Value, op: &str, name: &'static str) -> Result<&'a Value, PlanError> {
field(node, op, name)
}
fn seeds_field(node: &Value, op: &str) -> Result<Vec<(EntityId, f32)>, PlanError> {
let arr = field(node, op, "seeds")?.as_array().ok_or_else(|| PlanError::BadField {
op: op.to_string(),
field: "seeds",
wanted: "an array",
})?;
let mut out = Vec::with_capacity(arr.len());
for (index, seed) in arr.iter().enumerate() {
let bad = || PlanError::BadSeed { index, seed: seed.to_string() };
if let Some(s) = seed.as_str() {
out.push((EntityId::from_hex(s).ok_or_else(bad)?, 1.0));
continue;
}
let name = seed.get("name").and_then(Value::as_str).ok_or_else(bad)?;
let ty = seed.get("type").and_then(Value::as_str).ok_or_else(bad)?;
let conf = seed.get("confidence").and_then(Value::as_f64).unwrap_or(1.0) as f32;
out.push((EntityId::from_name_and_type(name, ty), conf));
}
Ok(out)
}