rig-core 0.44.0

//! Vector-search requests and composable metadata filters.
//!
//! ```
//! use rig_core::vector_store::request::{VectorSearchRequest, Filter, SearchFilter};
//!
//! let request: VectorSearchRequest = VectorSearchRequest::builder()
//!     .query("guide").samples(5)
//!     .filter(Filter::eq("category", serde_json::json!("manual"))).build();
//! assert_eq!(request.samples(), 5);
//! ```

use std::ops::Range;

use serde::{Deserialize, Serialize};

use crate::markers::{Missing, Provided};

/// A vector search request for querying a [`super::VectorStoreIndex`].
///
/// The type parameter `F` specifies the filter type (defaults to [`Filter<serde_json::Value>`]).
/// Use [`VectorSearchRequest::builder()`] to construct instances.
#[derive(Clone, Serialize, Deserialize, Debug)]
pub struct VectorSearchRequest<F = Filter<serde_json::Value>> {
    /// The query text to embed and search with.
    query: String,
    /// Maximum number of results to return.
    samples: u64,
    /// Minimum similarity score for results.
    threshold: Option<f64>,
    /// Filter expression to narrow results by metadata.
    filter: Option<F>,
}

impl<Filter> VectorSearchRequest<Filter> {
    /// Creates a [`VectorSearchRequestBuilder`] which you can use to instantiate this struct.
    pub fn builder() -> VectorSearchRequestBuilder<Filter> {
        VectorSearchRequestBuilder::<Filter>::default()
    }

    /// The query to be embedded and used in similarity search.
    pub fn query(&self) -> &str {
        &self.query
    }

    /// Returns the maximum number of results to return.
    pub fn samples(&self) -> u64 {
        self.samples
    }

    /// Returns the optional similarity threshold.
    pub fn threshold(&self) -> Option<f64> {
        self.threshold
    }

    /// Returns a reference to the optional filter expression.
    pub fn filter(&self) -> Option<&Filter> {
        self.filter.as_ref()
    }

    /// Transforms the filter type using the provided function.
    ///
    /// This is useful for converting between filter representations, such as
    /// translating the canonical [`super::request::Filter`] to a backend-specific filter type.
    pub fn map_filter<T, F>(self, f: F) -> VectorSearchRequest<T>
    where
        F: Fn(Filter) -> T,
    {
        VectorSearchRequest {
            query: self.query,
            samples: self.samples,
            threshold: self.threshold,
            filter: self.filter.map(f),
        }
    }

    /// Transforms the filter type using a provided function which can additionally return a result.
    ///
    /// Useful for converting between filter representations where the conversion can potentially fail (eg, unrepresentable or invalid values).
    pub fn try_map_filter<T, F>(self, f: F) -> Result<VectorSearchRequest<T>, FilterError>
    where
        F: Fn(Filter) -> Result<T, FilterError>,
    {
        let filter = self.filter.map(f).transpose()?;

        Ok(VectorSearchRequest {
            query: self.query,
            samples: self.samples,
            threshold: self.threshold,
            filter,
        })
    }
}

/// Errors from constructing or converting filter expressions.
#[derive(Debug, Clone, thiserror::Error)]
pub enum FilterError {
    #[error("Expected: {expected}, got: {got}")]
    Expected { expected: String, got: String },

    #[error("Cannot compile '{0}' to the backend's filter type")]
    TypeError(String),

    #[error("Missing field '{0}'")]
    MissingField(String),

    #[error("'{0}' must {1}")]
    Must(String, String),

    // NOTE: Uses String because `serde_json::Error` is not `Clone`.
    #[error("Filter serialization failed: {0}")]
    Serialization(String),
}

/// Trait for constructing filter expressions in vector search queries.
///
/// Uses [tagless final](https://nrinaudo.github.io/articles/tagless_final.html) encoding
/// for backend-agnostic filters. Use `SearchFilter::eq(...)` etc. directly and let
/// type inference resolve the concrete filter type.
pub trait SearchFilter {
    type Value;

    fn eq(key: impl AsRef<str>, value: Self::Value) -> Self;
    fn gt(key: impl AsRef<str>, value: Self::Value) -> Self;
    fn lt(key: impl AsRef<str>, value: Self::Value) -> Self;
    fn and(self, rhs: Self) -> Self;
    fn or(self, rhs: Self) -> Self;
}

/// A rendered SQL-style condition together with its positional bind parameters.
///
/// Parameters retain left-to-right placeholder order. Keys, operators, and
/// placeholders are interpolated verbatim; callers must validate or quote them
/// for the target backend. Only values are separated as bind parameters. The
/// condition records where each placeholder token sits, so
/// [`SqlCondition::render_placeholders`] rewrites placeholders without touching
/// spliced text that happens to contain the same characters.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct SqlCondition<P> {
    condition: String,
    params: Vec<P>,
    #[serde(default)]
    placeholders: Vec<Range<usize>>,
}

/// Hand-written so that `P` needs no [`Default`] of its own: a parameterless
/// empty condition is meaningful for every parameter type.
impl<P> Default for SqlCondition<P> {
    fn default() -> Self {
        Self::raw(String::new())
    }
}

impl<P> SqlCondition<P> {
    /// Renders `<key> <op> <placeholder>` bound to a single parameter, e.g.
    /// `price >= $`.
    pub fn binary(key: impl AsRef<str>, op: &str, placeholder: &str, value: P) -> Self {
        let mut this = Self::raw(format!("{} {op} ", key.as_ref()));
        this.push_placeholder(placeholder);
        this.params.push(value);
        this
    }

    /// Renders `<key> <op> (<placeholder>, ...)` with one placeholder per value,
    /// e.g. `id IN (?, ?)`.
    pub fn list(key: impl AsRef<str>, op: &str, placeholder: &str, values: Vec<P>) -> Self {
        let mut this = Self::raw(format!("{} {op} (", key.as_ref()));
        for i in 0..values.len() {
            if i > 0 {
                this.condition.push_str(", ");
            }
            this.push_placeholder(placeholder);
        }
        this.condition.push(')');
        this.params = values;
        this
    }

    /// Renders `<key> >= <placeholder> AND <key> <= <placeholder>` bound to
    /// `lo` then `hi`, e.g. `price >= ? AND price <= ?`.
    pub fn range(key: impl AsRef<str>, placeholder: &str, lo: P, hi: P) -> Self {
        let key = key.as_ref();
        let mut this = Self::raw(format!("{key} >= "));
        this.push_placeholder(placeholder);
        this.condition.push_str(&format!(" AND {key} <= "));
        this.push_placeholder(placeholder);
        this.params = vec![lo, hi];
        this
    }

    /// Renders `<key> between <placeholder> and <placeholder>` bound to `lo`
    /// then `hi`, e.g. `price between ? and ?`.
    pub fn between(key: impl AsRef<str>, placeholder: &str, lo: P, hi: P) -> Self {
        let mut this = Self::raw(format!("{} between ", key.as_ref()));
        this.push_placeholder(placeholder);
        this.condition.push_str(" and ");
        this.push_placeholder(placeholder);
        this.params = vec![lo, hi];
        this
    }

    /// Wraps an already-rendered, parameterless condition such as `id is null`.
    /// The text holds no placeholders, whatever characters it contains.
    pub fn raw(condition: impl Into<String>) -> Self {
        Self {
            condition: condition.into(),
            params: Vec::new(),
            placeholders: Vec::new(),
        }
    }

    /// Conjoins two conditions as `(lhs) AND (rhs)`, concatenating their parameters.
    pub fn and(self, rhs: Self) -> Self {
        self.combine("AND", rhs)
    }

    /// Disjoins two conditions as `(lhs) OR (rhs)`, concatenating their parameters.
    pub fn or(self, rhs: Self) -> Self {
        self.combine("OR", rhs)
    }

    /// Negates the condition as `NOT (condition)`, keeping its parameters.
    pub fn not(self) -> Self {
        let mut this = Self::raw("NOT (");
        this.append(self);
        this.condition.push(')');
        this
    }

    fn combine(self, joiner: &str, rhs: Self) -> Self {
        let mut this = Self::raw("(");
        this.append(self);
        this.condition.push_str(&format!(") {joiner} ("));
        this.append(rhs);
        this.condition.push(')');
        this
    }

    fn push_placeholder(&mut self, placeholder: &str) {
        let start = self.condition.len();
        self.condition.push_str(placeholder);
        self.placeholders.push(start..self.condition.len());
    }

    fn append(&mut self, other: Self) {
        let offset = self.condition.len();
        self.condition.push_str(&other.condition);
        self.params.extend(other.params);
        self.placeholders.extend(
            other
                .placeholders
                .into_iter()
                .map(|range| range.start + offset..range.end + offset),
        );
    }

    /// The rendered condition, with placeholders as the caller supplied them.
    pub fn condition(&self) -> &str {
        &self.condition
    }

    /// The bind parameters, in placeholder order.
    pub fn params(&self) -> &[P] {
        &self.params
    }

    /// Returns the condition with the placeholder for parameter `i` (zero-based,
    /// in [`SqlCondition::params`] order) replaced by `placeholder(i)`, e.g.
    /// `|i| format!("${}", i + 1)` for Postgres. Only tokens emitted by the
    /// constructors are replaced; keys and [`SqlCondition::raw`] text are
    /// copied verbatim.
    pub fn render_placeholders<D: std::fmt::Display>(
        &self,
        mut placeholder: impl FnMut(usize) -> D,
    ) -> String {
        use std::fmt::Write as _;

        let mut out = String::with_capacity(self.condition.len() + 2 * self.placeholders.len());
        let mut copied = 0;
        for (i, range) in self.placeholders.iter().enumerate() {
            // Ranges only come from the constructors or deserialization; skip
            // any that do not fit the text rather than panic.
            let (Some(before), Some(_)) = (
                self.condition.get(copied..range.start),
                self.condition.get(range.clone()),
            ) else {
                continue;
            };
            out.push_str(before);
            let _ = write!(out, "{}", placeholder(i));
            copied = range.end;
        }
        out.push_str(self.condition.get(copied..).unwrap_or_default());
        out
    }

    /// Consumes the condition, returning the rendered text and its parameters.
    pub fn into_parts(self) -> (String, Vec<P>) {
        (self.condition, self.params)
    }
}

/// Converts the canonical JSON-valued [`Filter`] used by type-erased vector
/// searches into a backend's native filter representation.
///
/// JSON-valued [`SearchFilter`] implementations receive this automatically.
/// Backends with native value types implement the conversion once here rather
/// than hand-writing both the bus's `RetrieveAdapter`
/// methods.
pub trait DynamicSearchFilter: SearchFilter + Sized {
    /// Converts a canonical dynamic filter into this backend's filter type.
    fn from_dynamic_filter(filter: Filter<serde_json::Value>) -> Result<Self, FilterError>;

    /// Normalizes a document returned through the type-erased search surface.
    ///
    /// The default retains documents unchanged. The JSON-valued blanket
    /// implementation recursively removes arrays longer than 400 elements;
    /// a removed root array becomes null.
    fn normalize_dynamic_document(document: serde_json::Value) -> serde_json::Value {
        document
    }
}

impl<F> DynamicSearchFilter for F
where
    F: SearchFilter<Value = serde_json::Value>,
{
    fn from_dynamic_filter(filter: Filter<serde_json::Value>) -> Result<Self, FilterError> {
        Ok(filter.interpret())
    }

    fn normalize_dynamic_document(document: serde_json::Value) -> serde_json::Value {
        prune_document(document).unwrap_or_default()
    }
}

fn prune_document(document: serde_json::Value) -> Option<serde_json::Value> {
    match document {
        serde_json::Value::Object(mut map) => {
            let new_map = map
                .iter_mut()
                .filter_map(|(key, value)| {
                    prune_document(value.take()).map(|value| (key.clone(), value))
                })
                .collect::<serde_json::Map<_, _>>();

            Some(serde_json::Value::Object(new_map))
        }
        serde_json::Value::Array(vec) if vec.len() > 400 => None,
        serde_json::Value::Array(vec) => Some(serde_json::Value::Array(
            vec.into_iter().filter_map(prune_document).collect(),
        )),
        value => Some(value),
    }
}

/// Canonical, serializable filter representation.
///
/// Use for serialization, runtime inspection, or translating between backends via
/// [`Filter::interpret`]. Prefer [`SearchFilter`] trait methods for writing queries.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Filter<V>
where
    V: std::fmt::Debug + Clone,
{
    Eq(String, V),
    Gt(String, V),
    Lt(String, V),
    And(Box<Self>, Box<Self>),
    Or(Box<Self>, Box<Self>),
}

impl<V> SearchFilter for Filter<V>
where
    V: std::fmt::Debug + Clone + Serialize + serde::de::DeserializeOwned,
{
    type Value = V;

    /// Select values where the entry at `key` is equal to `value`
    fn eq(key: impl AsRef<str>, value: Self::Value) -> Self {
        Self::Eq(key.as_ref().to_owned(), value)
    }

    /// Select values where the entry at `key` is greater than `value`
    fn gt(key: impl AsRef<str>, value: Self::Value) -> Self {
        Self::Gt(key.as_ref().to_owned(), value)
    }

    /// Select values where the entry at `key` is less than `value`
    fn lt(key: impl AsRef<str>, value: Self::Value) -> Self {
        Self::Lt(key.as_ref().to_owned(), value)
    }

    /// Select values where the entry satisfies `self` *and* `rhs`
    fn and(self, rhs: Self) -> Self {
        Self::And(self.into(), rhs.into())
    }

    /// Select values where the entry satisfies `self` *or* `rhs`
    fn or(self, rhs: Self) -> Self {
        Self::Or(self.into(), rhs.into())
    }
}

impl<V> Filter<V>
where
    V: std::fmt::Debug + Clone,
{
    /// Converts this filter into a backend-specific filter type.
    pub fn interpret<F>(self) -> F
    where
        F: SearchFilter<Value = V>,
    {
        self.interpret_with(|v| v)
    }

    /// Converts this filter into a backend-specific filter type, converting
    /// each leaf value with `conv`.
    pub fn interpret_with<F, W>(self, conv: impl Fn(V) -> W + Copy) -> F
    where
        F: SearchFilter<Value = W>,
    {
        match self.try_interpret(|v| Ok::<W, std::convert::Infallible>(conv(v))) {
            Ok(filter) => filter,
            Err(never) => match never {},
        }
    }

    /// Converts this filter into a backend-specific filter type, converting
    /// each leaf value with `conv`. Fails on the first value that fails to convert.
    pub fn try_interpret<F, W, E>(self, conv: impl Fn(V) -> Result<W, E> + Copy) -> Result<F, E>
    where
        F: SearchFilter<Value = W>,
    {
        Ok(match self {
            Self::Eq(key, val) => F::eq(key, conv(val)?),
            Self::Gt(key, val) => F::gt(key, conv(val)?),
            Self::Lt(key, val) => F::lt(key, conv(val)?),
            Self::And(lhs, rhs) => F::and(lhs.try_interpret(conv)?, rhs.try_interpret(conv)?),
            Self::Or(lhs, rhs) => F::or(lhs.try_interpret(conv)?, rhs.try_interpret(conv)?),
        })
    }
}

impl Filter<serde_json::Value> {
    /// Tests whether a JSON document satisfies this filter.
    ///
    /// Looks up top-level object fields. Missing fields fail all leaves.
    /// Equality accepts numeric comparison or structural JSON equality; ordering
    /// compares only number, string, boolean, or null pairs of the same kind.
    /// `And` and `Or` short-circuit.
    pub fn satisfies(&self, value: &serde_json::Value) -> bool {
        use Filter::*;
        use serde_json::{Value, Value::*};
        use std::cmp::Ordering;

        fn compare_pair(l: &Value, r: &Value) -> Option<Ordering> {
            match (l, r) {
                // Prefer shared integer representations to avoid losing precision
                // beyond 2^53; other numeric pairs fall back to f64.
                (Number(l), Number(r)) => {
                    if let (Some(l), Some(r)) = (l.as_i64(), r.as_i64()) {
                        Some(l.cmp(&r))
                    } else if let (Some(l), Some(r)) = (l.as_u64(), r.as_u64()) {
                        Some(l.cmp(&r))
                    } else {
                        l.as_f64()
                            .zip(r.as_f64())
                            .and_then(|(l, r)| l.partial_cmp(&r))
                    }
                }
                (String(l), String(r)) => Some(l.cmp(r)),
                (Null, Null) => Some(Ordering::Equal),
                (Bool(l), Bool(r)) => Some(l.cmp(r)),
                _ => None,
            }
        }

        match self {
            // Numbers compare numerically so `5` matches `5.0`, consistent with
            // `Gt`/`Lt`; other JSON types fall back to structural equality so
            // strings/bools/arrays/objects still match exactly.
            Eq(k, v) => value
                .get(k)
                .is_some_and(|field| compare_pair(field, v) == Some(Ordering::Equal) || field == v),
            Gt(k, v) => value
                .get(k)
                .and_then(|field| compare_pair(field, v))
                .is_some_and(|ord| ord == Ordering::Greater),
            Lt(k, v) => value
                .get(k)
                .and_then(|field| compare_pair(field, v))
                .is_some_and(|ord| ord == Ordering::Less),
            And(l, r) => l.satisfies(value) && r.satisfies(value),
            Or(l, r) => l.satisfies(value) || r.satisfies(value),
        }
    }
}

/// Builder for [`VectorSearchRequest`]. Requires `query` and `samples`.
#[derive(Clone, Serialize, Deserialize, Debug)]
pub struct VectorSearchRequestBuilder<F = Filter<serde_json::Value>, Q = Missing, S = Missing> {
    query: Q,
    samples: S,
    threshold: Option<f64>,
    filter: Option<F>,
}

impl<F> Default for VectorSearchRequestBuilder<F, Missing, Missing> {
    fn default() -> Self {
        Self {
            query: Missing,
            samples: Missing,
            threshold: None,
            filter: None,
        }
    }
}

impl<F, Q, S> VectorSearchRequestBuilder<F, Q, S>
where
    F: SearchFilter,
{
    /// Sets the query text. Required.
    pub fn query<T>(self, query: T) -> VectorSearchRequestBuilder<F, Provided<String>, S>
    where
        T: Into<String>,
    {
        VectorSearchRequestBuilder {
            query: Provided(query.into()),
            samples: self.samples,
            threshold: self.threshold,
            filter: self.filter,
        }
    }

    /// Sets the maximum number of results. Required.
    pub fn samples(self, samples: u64) -> VectorSearchRequestBuilder<F, Q, Provided<u64>> {
        VectorSearchRequestBuilder {
            query: self.query,
            samples: Provided(samples),
            threshold: self.threshold,
            filter: self.filter,
        }
    }

    /// Sets the minimum similarity threshold.
    pub fn threshold(mut self, threshold: f64) -> Self {
        self.threshold = Some(threshold);
        self
    }

    /// Sets a filter expression.
    pub fn filter(mut self, filter: F) -> Self {
        self.filter = Some(filter);
        self
    }
}

impl<F> VectorSearchRequestBuilder<F, Provided<String>, Provided<u64>> {
    /// Builds without validating query emptiness, sample count, or threshold.
    pub fn build(self) -> VectorSearchRequest<F> {
        VectorSearchRequest {
            query: self.query.0,
            samples: self.samples.0,
            threshold: self.threshold,
            filter: self.filter,
        }
    }
}

#[cfg(test)]
mod tests;