use lora_store::{
cosine_similarity_bounded, cosine_similarity_raw, dot_product, euclidean_distance,
euclidean_distance_squared, euclidean_norm, euclidean_similarity, hamming_distance,
manhattan_distance, manhattan_norm, LoraVector, RawCoordinate, VectorCoordinateType,
};
use crate::value::LoraValue;
use super::errors::set_eval_error;
fn coerce_list_to_raw_coords(items: &[LoraValue]) -> Result<Vec<RawCoordinate>, String> {
let mut out = Vec::with_capacity(items.len());
for item in items {
match item {
LoraValue::Int(i) => out.push(RawCoordinate::Int(*i)),
LoraValue::Float(f) => {
if !f.is_finite() {
return Err("vector coordinates cannot be NaN or Infinity".to_string());
}
out.push(RawCoordinate::Float(*f));
}
LoraValue::List(_) => {
return Err("vector coordinates cannot contain nested lists".to_string());
}
LoraValue::Null => {
return Err("vector coordinates cannot be null".to_string());
}
other => {
return Err(format!(
"vector coordinates must be numeric, got `{}`",
crate::errors::value_kind(other)
));
}
}
}
Ok(out)
}
fn coerce_similarity_input(value: &LoraValue) -> Result<Option<LoraVector>, String> {
match value {
LoraValue::Null => Ok(None),
LoraValue::Vector(v) => Ok(Some(v.clone())),
LoraValue::List(items) => {
let raw = coerce_list_to_raw_coords(items)?;
let dim = raw.len() as i64;
if dim == 0 {
return Err("vector similarity cannot be computed on an empty list".to_string());
}
LoraVector::try_new(raw, dim, VectorCoordinateType::Float32)
.map(Some)
.map_err(|e| e.to_string())
}
other => Err(format!(
"expected VECTOR or LIST<NUMBER>, got `{}`",
crate::errors::value_kind(other)
)),
}
}
pub(super) fn eval_vector_sim_cosine(args: &[LoraValue]) -> LoraValue {
let a = args.first();
let b = args.get(1);
let (Some(a), Some(b)) = (a, b) else {
return LoraValue::Null;
};
let av = match coerce_similarity_input(a) {
Ok(Some(v)) => v,
Ok(None) => return LoraValue::Null,
Err(e) => {
set_eval_error(e);
return LoraValue::Null;
}
};
let bv = match coerce_similarity_input(b) {
Ok(Some(v)) => v,
Ok(None) => return LoraValue::Null,
Err(e) => {
set_eval_error(e);
return LoraValue::Null;
}
};
if av.dimension != bv.dimension {
set_eval_error(format!(
"`vector.similarity()` requires equal dimensions, got {} and {}",
av.dimension, bv.dimension
));
return LoraValue::Null;
}
match cosine_similarity_bounded(&av, &bv) {
Some(s) => LoraValue::Float(s),
None => LoraValue::Null,
}
}
pub(super) fn eval_vector_sim_euclidean(args: &[LoraValue]) -> LoraValue {
let a = args.first();
let b = args.get(1);
let (Some(a), Some(b)) = (a, b) else {
return LoraValue::Null;
};
let av = match coerce_similarity_input(a) {
Ok(Some(v)) => v,
Ok(None) => return LoraValue::Null,
Err(e) => {
set_eval_error(e);
return LoraValue::Null;
}
};
let bv = match coerce_similarity_input(b) {
Ok(Some(v)) => v,
Ok(None) => return LoraValue::Null,
Err(e) => {
set_eval_error(e);
return LoraValue::Null;
}
};
if av.dimension != bv.dimension {
set_eval_error(format!(
"`vector.similarity()` requires equal dimensions, got {} and {}",
av.dimension, bv.dimension
));
return LoraValue::Null;
}
match euclidean_similarity(&av, &bv) {
Some(s) => LoraValue::Float(s),
None => LoraValue::Null,
}
}
pub(super) fn eval_vector_distance_fn(args: &[LoraValue]) -> LoraValue {
let a = args.first();
let b = args.get(1);
let metric = args.get(2);
if matches!(a, Some(LoraValue::Null)) || matches!(b, Some(LoraValue::Null)) {
return LoraValue::Null;
}
let (Some(LoraValue::Vector(av)), Some(LoraValue::Vector(bv))) = (a, b) else {
set_eval_error("`vector.distance()` requires two VECTOR arguments".to_string());
return LoraValue::Null;
};
if av.dimension != bv.dimension {
set_eval_error(format!(
"`vector.distance()` requires equal dimensions, got {} and {}",
av.dimension, bv.dimension
));
return LoraValue::Null;
}
let metric_str = match metric {
Some(LoraValue::String(s)) => s.clone(),
Some(LoraValue::Null) => return LoraValue::Null,
_ => {
set_eval_error("`vector.distance()` metric must be a string/identifier".to_string());
return LoraValue::Null;
}
};
let result = match metric_str.to_ascii_uppercase().as_str() {
"EUCLIDEAN" => euclidean_distance(av, bv),
"EUCLIDEAN_SQUARED" => euclidean_distance_squared(av, bv),
"MANHATTAN" => manhattan_distance(av, bv),
"COSINE" => cosine_similarity_raw(av, bv).map(|s| 1.0 - s),
"DOT" => dot_product(av, bv).map(|d| -d),
"HAMMING" => hamming_distance(av, bv),
other => {
set_eval_error(format!("unknown vector distance metric `{other}`"));
return LoraValue::Null;
}
};
result.map(LoraValue::Float).unwrap_or(LoraValue::Null)
}
pub(super) fn eval_vector_norm_fn(args: &[LoraValue]) -> LoraValue {
let v = args.first();
let metric = args.get(1);
if matches!(v, Some(LoraValue::Null)) {
return LoraValue::Null;
}
let Some(LoraValue::Vector(v)) = v else {
set_eval_error("`vector.norm()` requires a VECTOR argument".to_string());
return LoraValue::Null;
};
let metric_str = match metric {
Some(LoraValue::String(s)) => s.clone(),
Some(LoraValue::Null) => return LoraValue::Null,
_ => {
set_eval_error("`vector.norm()` metric must be a string/identifier".to_string());
return LoraValue::Null;
}
};
match metric_str.to_ascii_uppercase().as_str() {
"EUCLIDEAN" => LoraValue::Float(euclidean_norm(v)),
"MANHATTAN" => LoraValue::Float(manhattan_norm(v)),
other => {
set_eval_error(format!("unknown vector norm metric `{other}`"));
LoraValue::Null
}
}
}