swiftide_core/

node.rs

//! This module defines the `Node` struct and its associated methods.
//!
//! `Node` represents a unit of data in the indexing process, containing metadata,
//! the data chunk itself, and an optional vector representation.
//!
//! # Overview
//!
//! The `Node` struct is designed to encapsulate all necessary information for a single
//! unit of data being processed in the indexing pipeline. It includes fields for an identifier,
//! file path, data chunk, optional vector representation, and metadata.
//!
//! The struct provides methods to convert the node into an embeddable string format and to
//! calculate a hash value for the node based on its path and chunk.
//!
//! # Usage
//!
//! The `Node` struct is used throughout the indexing pipeline to represent and process
//! individual units of data. It is particularly useful in scenarios where metadata and data chunks
//! need to be processed together.
use std::{
    collections::HashMap,
    fmt::Debug,
    hash::{Hash, Hasher},
    os::unix::ffi::OsStrExt,
    path::PathBuf,
};

use derive_builder::Builder;
use itertools::Itertools;
use serde::{Deserialize, Serialize};

use crate::{Embedding, SparseEmbedding, metadata::Metadata};

/// Helper trait for types that can be used as data chunks in a `Node`.
/// For now always expects an owned value
///
/// A chunk must be able to yield its bytes, be cloned (not while streaming), and be sent across
/// threads.
pub trait Chunk: Clone + Send + Sync + Debug + AsRef<[u8]> + 'static {}
impl<T> Chunk for T where T: Clone + Send + Sync + Debug + AsRef<[u8]> + 'static {}

/// Represents a unit of data in the indexing process.
///
/// `Node` encapsulates all necessary information for a single unit of data being processed
/// in the indexing pipeline. It includes fields for an identifier, file path, data chunk, optional
/// vector representation, and metadata.
#[derive(Default, Clone, Serialize, Deserialize, PartialEq, Builder)]
#[builder(setter(into, strip_option), build_fn(error = "anyhow::Error"))]
pub struct Node<T: Chunk> {
    /// File path associated with the node.
    #[builder(default)]
    pub path: PathBuf,
    /// Data chunk contained in the node.
    pub chunk: T,
    /// Optional vector representation of embedded data.
    #[builder(default)]
    pub vectors: Option<HashMap<EmbeddedField, Embedding>>,
    /// Optional sparse vector representation of embedded data.
    #[builder(default)]
    pub sparse_vectors: Option<HashMap<EmbeddedField, SparseEmbedding>>,
    /// Metadata associated with the node.
    #[builder(default)]
    pub metadata: Metadata,
    /// Mode of embedding data Chunk and Metadata
    #[builder(default)]
    pub embed_mode: EmbedMode,
    /// Size of the input this node was originally derived from in bytes
    #[builder(default)]
    pub original_size: usize,
    /// Offset of the chunk relative to the start of the input this node was originally derived
    /// from in bytes
    #[builder(default)]
    pub offset: usize,
    /// Id of the first node this node was derived from. Set when a node is created via
    /// [`Node::build_from_other`] (chunking, conversions) so derived nodes keep pointing at the
    /// node that entered the pipeline.
    #[builder(default)]
    pub parent_id: Option<uuid::Uuid>,
}

pub type TextNode = Node<String>;

impl<T: Chunk> NodeBuilder<T> {
    pub fn maybe_sparse_vectors(
        &mut self,
        sparse_vectors: Option<HashMap<EmbeddedField, SparseEmbedding>>,
    ) -> &mut Self {
        self.sparse_vectors = Some(sparse_vectors);
        self
    }

    pub fn maybe_vectors(
        &mut self,
        vectors: Option<HashMap<EmbeddedField, Embedding>>,
    ) -> &mut Self {
        self.vectors = Some(vectors);
        self
    }
}

impl<T: Chunk> Debug for Node<T> {
    /// Formats the node for debugging purposes.
    ///
    /// This method is used to provide a human-readable representation of the node when debugging.
    /// The vector field is displayed as the number of elements in the vector if present.
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Node")
            .field("id", &self.id())
            .field("path", &self.path)
            .field("chunk", &self.chunk)
            .field("metadata", &self.metadata)
            .field(
                "vectors",
                &self
                    .vectors
                    .iter()
                    .flat_map(HashMap::iter)
                    .map(|(embed_type, vec)| format!("'{embed_type}': {}", vec.len()))
                    .join(","),
            )
            .field(
                "sparse_vectors",
                &self
                    .sparse_vectors
                    .iter()
                    .flat_map(HashMap::iter)
                    .map(|(embed_type, vec)| {
                        format!(
                            "'{embed_type}': indices({}), values({})",
                            vec.indices.len(),
                            vec.values.len()
                        )
                    })
                    .join(","),
            )
            .field("embed_mode", &self.embed_mode)
            .finish()
    }
}

impl<T: Chunk> Node<T> {
    /// Builds a new instance of `Node`, returning a `NodeBuilder`. Copies
    /// over the fields from the provided `Node`.
    ///
    /// The new node points at the node that entered the pipeline via [`Node::parent_id`].
    pub fn build_from_other(node: &Node<T>) -> NodeBuilder<T> {
        NodeBuilder::default()
            .path(node.path.clone())
            .chunk(node.chunk.clone())
            .metadata(node.metadata.clone())
            .maybe_vectors(node.vectors.clone())
            .maybe_sparse_vectors(node.sparse_vectors.clone())
            .embed_mode(node.embed_mode)
            .original_size(node.original_size)
            .offset(node.offset)
            .parent_id(node.parent_id.unwrap_or_else(|| node.id()))
            .to_owned()
    }

    /// Creates a new instance of `NodeBuilder.`
    pub fn builder<VALUE: Chunk + Clone>() -> NodeBuilder<VALUE> {
        NodeBuilder::default()
    }

    /// Creates a new instance of `Node` with the specified data chunk.
    ///
    /// The other fields are set to their default values.
    pub fn new(chunk: impl Into<String>) -> Node<String> {
        let chunk = chunk.into();
        let original_size = chunk.len();
        Node {
            chunk,
            original_size,
            ..Default::default()
        }
    }

    pub fn with_metadata(&mut self, metadata: impl Into<Metadata>) -> &mut Self {
        self.metadata = metadata.into();
        self
    }

    pub fn with_vectors(
        &mut self,
        vectors: impl Into<HashMap<EmbeddedField, Embedding>>,
    ) -> &mut Self {
        self.vectors = Some(vectors.into());
        self
    }

    pub fn with_sparse_vectors(
        &mut self,
        sparse_vectors: impl Into<HashMap<EmbeddedField, SparseEmbedding>>,
    ) -> &mut Self {
        self.sparse_vectors = Some(sparse_vectors.into());
        self
    }

    /// Retrieve the identifier of the node.
    ///
    /// Calculates the identifier of the node based on its path and chunk as bytes, returning a
    /// UUID (v3).
    ///
    /// WARN: Does not memoize the id. Use sparingly.
    pub fn id(&self) -> uuid::Uuid {
        // Calculate the identifier based on the path and chunk as bytes
        let bytes = [self.path.as_os_str().as_bytes(), self.chunk.as_ref()].concat();

        uuid::Uuid::new_v3(&uuid::Uuid::NAMESPACE_OID, &bytes)
    }

    /// Returns the id of the first node this node was derived from, if any.
    pub fn parent_id(&self) -> Option<uuid::Uuid> {
        self.parent_id
    }
}

impl Node<String> {
    /// Derives nodes from chunks while retaining the first ancestor as their parent.
    ///
    /// Resolves the parent before removing the input text, then reuses the remaining
    /// fields for every output without cloning the full input text.
    pub fn into_chunks(
        self,
        chunks: impl IntoIterator<Item = String>,
    ) -> impl Iterator<Item = Self> {
        let parent_id = self.parent_id.unwrap_or_else(|| self.id());
        let mut template = self;
        template.chunk.clear();
        template.parent_id = Some(parent_id);

        chunks.into_iter().map(move |chunk| {
            let mut node = template.clone();
            node.chunk = chunk;
            node
        })
    }

    /// Creates embeddable data depending on chosen `EmbedMode`.
    ///
    /// # Returns
    ///
    /// Embeddable data mapped to their `EmbeddedField`.
    pub fn as_embeddables(&self) -> Vec<(EmbeddedField, String)> {
        // TODO: Cow and borrow the inner data + generic
        let mut embeddables = Vec::new();

        if self.embed_mode == EmbedMode::SingleWithMetadata || self.embed_mode == EmbedMode::Both {
            embeddables.push((EmbeddedField::Combined, self.combine_chunk_with_metadata()));
        }

        if self.embed_mode == EmbedMode::PerField || self.embed_mode == EmbedMode::Both {
            embeddables.push((EmbeddedField::Chunk, self.chunk.clone()));
            for (name, value) in &self.metadata {
                let value = value
                    .as_str()
                    .map_or_else(|| value.to_string(), ToString::to_string);
                embeddables.push((EmbeddedField::Metadata(name.clone()), value));
            }
        }

        embeddables
    }

    /// Converts the node into an [`self::EmbeddedField::Combined`] type of embeddable.
    ///
    /// This embeddable format consists of the metadata formatted as key-value pairs, each on a new
    /// line, followed by the data chunk.
    ///
    /// # Returns
    ///
    /// A string representing the embeddable format of the node.
    fn combine_chunk_with_metadata(&self) -> String {
        // Metadata formatted by newlines joined with the chunk
        let metadata = self
            .metadata
            .iter()
            .map(|(k, v)| {
                let v = v
                    .as_str()
                    .map_or_else(|| v.to_string(), ToString::to_string);

                format!("{k}: {v}")
            })
            .collect::<Vec<String>>()
            .join("\n");

        format!("{}\n{}", metadata, self.chunk)
    }
}

impl Hash for Node<String> {
    /// Hashes the node based on its path and chunk.
    ///
    /// This method is used by the `calculate_hash` method to generate a hash value for the node.
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.path.hash(state);
        self.chunk.hash(state);
    }
}

impl<T: Into<String>> From<T> for Node<String> {
    fn from(value: T) -> Self {
        let value: String = value.into();
        Node::<String>::new(value)
    }
}

/// Embed mode of the pipeline.
#[derive(Copy, Debug, Default, Clone, Serialize, Deserialize, PartialEq)]
pub enum EmbedMode {
    #[default]
    /// Embedding Chunk of data combined with Metadata.
    SingleWithMetadata,
    /// Embedding Chunk of data and every Metadata separately.
    PerField,
    /// Embedding Chunk of data and every Metadata separately and Chunk of data combined with
    /// Metadata.
    Both,
}

/// Type of Embeddable stored in model.
#[derive(
    Clone, Default, Serialize, Deserialize, PartialEq, Eq, Hash, strum_macros::Display, Debug,
)]
pub enum EmbeddedField {
    #[default]
    /// Embeddable created from Chunk of data combined with Metadata.
    Combined,
    /// Embeddable created from Chunk of data only.
    Chunk,
    /// Embeddable created from Metadata.
    /// String stores Metadata name.
    #[strum(to_string = "Metadata: {0}")]
    Metadata(String),
}

impl EmbeddedField {
    /// Returns the name of the field when it would be a sparse vector
    pub fn sparse_field_name(&self) -> String {
        format!("{self}_sparse")
    }

    /// Returns the name of the field when it would be a dense vector
    pub fn field_name(&self) -> String {
        format!("{self}")
    }
}

#[allow(clippy::from_over_into)]
impl Into<String> for EmbeddedField {
    fn into(self) -> String {
        self.to_string()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use test_case::test_case;

    #[test_case(&EmbeddedField::Combined, ["Combined", "Combined_sparse"])]
    #[test_case(&EmbeddedField::Chunk, ["Chunk", "Chunk_sparse"])]
    #[test_case(&EmbeddedField::Metadata("test".into()), ["Metadata: test", "Metadata: test_sparse"])]
    fn field_name_tests(embedded_field: &EmbeddedField, expected: [&str; 2]) {
        assert_eq!(embedded_field.field_name(), expected[0]);
        assert_eq!(embedded_field.sparse_field_name(), expected[1]);
    }

    #[test]
    fn test_debugging_node_with_utf8_char_boundary() {
        let node = Node::from("🦀".repeat(101));
        // Single char
        let _ = format!("{node:?}");

        // With invalid char boundary
        let node = Node::from("Jürgen".repeat(100));
        let _ = format!("{node:?}");
    }

    #[test]
    fn test_build_from_other_without_vectors() {
        let original_node = Node::from("test_chunk")
            .with_metadata(Metadata::default())
            .with_vectors(HashMap::new())
            .with_sparse_vectors(HashMap::new())
            .to_owned();

        let builder = Node::build_from_other(&original_node);
        let new_node = builder.build().unwrap();

        assert_eq!(new_node.parent_id(), Some(original_node.id()));
        assert_eq!(new_node.chunk, original_node.chunk);
        assert_eq!(new_node.path, original_node.path);
        assert_eq!(new_node.metadata, original_node.metadata);
        assert_eq!(new_node.vectors, original_node.vectors);
        assert_eq!(new_node.sparse_vectors, original_node.sparse_vectors);
    }

    #[test]
    fn test_build_from_other_with_vectors() {
        let mut vectors = HashMap::new();
        vectors.insert(EmbeddedField::Chunk, Embedding::default());

        let mut sparse_vectors = HashMap::new();
        sparse_vectors.insert(
            EmbeddedField::Chunk,
            SparseEmbedding {
                indices: vec![],
                values: vec![],
            },
        );

        let original_node = Node::from("test_chunk")
            .with_metadata(Metadata::default())
            .with_vectors(vectors.clone())
            .with_sparse_vectors(sparse_vectors.clone())
            .to_owned();

        let builder = Node::build_from_other(&original_node);
        let new_node = builder.build().unwrap();

        assert_eq!(new_node.parent_id(), Some(original_node.id()));
        assert_eq!(new_node.chunk, original_node.chunk);
        assert_eq!(new_node.vectors, original_node.vectors);
        assert_eq!(new_node.sparse_vectors, original_node.sparse_vectors);
    }

    #[test]
    fn into_chunks_preserves_fields_and_first_ancestor() {
        let source = TextNode::builder()
            .path("fixtures/日本語.txt")
            .chunk("αβγ parent".to_string())
            .metadata(Metadata::from([("language", "日本語")]))
            .vectors(HashMap::from([(EmbeddedField::Chunk, vec![1.0, 2.0])]))
            .sparse_vectors(HashMap::from([(
                EmbeddedField::Chunk,
                SparseEmbedding {
                    indices: vec![2],
                    values: vec![0.5],
                },
            )]))
            .embed_mode(EmbedMode::Both)
            .original_size(100usize)
            .offset(7usize)
            .build()
            .unwrap();
        let parent_id = source.id();
        let chunks = vec!["αβγ".to_string(), "日本語".to_string()];
        let outputs: Vec<_> = source.clone().into_chunks(chunks.clone()).collect();
        assert_eq!(outputs.len(), 2);
        for (output, chunk) in outputs.iter().zip(chunks) {
            assert_eq!(output.chunk, chunk);
            assert_eq!(output.path, source.path);
            assert_eq!(output.metadata, source.metadata);
            assert_eq!(output.vectors, source.vectors);
            assert_eq!(output.sparse_vectors, source.sparse_vectors);
            assert_eq!(output.embed_mode, source.embed_mode);
            assert_eq!(output.original_size, source.original_size);
            assert_eq!(output.offset, source.offset);
            assert_eq!(output.parent_id, Some(parent_id));
        }

        let descendant: Vec<_> = outputs[0]
            .clone()
            .into_chunks(["child".to_string()])
            .collect();
        assert_eq!(descendant[0].parent_id, Some(parent_id));

        let no_outputs: Vec<_> = source.into_chunks(Vec::<String>::new()).collect();
        assert!(no_outputs.is_empty());
    }

    #[test]
    fn into_chunks_preserves_explicit_parent_and_live_id() {
        let mut source = Node::from("original");
        let original_id = source.id();
        source.path = "changed.txt".into();
        source.chunk = "changed".into();
        assert_ne!(source.id(), original_id);

        let explicit_parent = uuid::Uuid::new_v4();
        source.parent_id = Some(explicit_parent);
        let children: Vec<_> = source.into_chunks(["one".into(), "two".into()]).collect();
        assert_eq!(children.len(), 2);
        assert!(
            children
                .iter()
                .all(|child| child.parent_id == Some(explicit_parent))
        );
    }

    #[test]
    fn test_build_from_other_keeps_first_ancestor_as_parent() {
        let original_node = Node::from("test_chunk");
        let child_node = Node::build_from_other(&original_node).build().unwrap();
        let grandchild_node = Node::build_from_other(&child_node).build().unwrap();

        assert_eq!(original_node.parent_id(), None);
        assert_eq!(child_node.parent_id(), Some(original_node.id()));
        assert_eq!(grandchild_node.parent_id(), Some(original_node.id()));
    }
}