kiromi_ai_memory/

graph.rs

// SPDX-License-Identifier: Apache-2.0 OR MIT
//! Plan 10 multi-hop traversal types — `NodeRef`, `EdgeKind`, `Graph`,
//! `TraverseOpts`. The actual `Memory::traverse` implementation lives on
//! [`crate::Memory`] (see `handle.rs`).
//!
//! `Graph` is purposefully simple — `Vec<GraphNode>` and `Vec<GraphEdge>` —
//! so callers can render it without a graph dependency. Plan 10 keeps this
//! Rust-only; Swift consumers go through `Memory::traverse_json` to get
//! `serde_json::to_string(&graph)`.

use std::collections::BTreeSet;

use serde::{Deserialize, Serialize};

use crate::memory::MemoryRef;
use crate::partition::PartitionPath;
use crate::summary::{SummaryRef, SummarySubject};

/// One node in a traversal graph.
///
/// Three kinds: a memory, a summary, or a partition. Memory and Summary
/// carry the full ref; partitions carry just the path. Future fields
/// (cached labels, attrs) will land on `GraphNode` rather than here.
#[non_exhaustive]
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(tag = "kind", content = "value", rename_all = "snake_case")]
pub enum NodeRef {
    /// A memory.
    Memory(MemoryRef),
    /// A summary.
    Summary(SummaryRef),
    /// A partition. The tenant root is represented as the
    /// [`crate::partition::tenant_root_path`] sentinel.
    Partition(PartitionPath),
}

/// One edge in a traversal graph. `EdgeKind` distinguishes the relation —
/// link, summary input, parent partition, partition contains.
#[non_exhaustive]
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(tag = "kind", content = "value", rename_all = "snake_case")]
pub enum EdgeKind {
    /// `memory ↔ memory` explicit link.
    Link,
    /// `summary → input` — the summary cites this input. The variant
    /// records the summary's subject so renderers can group "the summary
    /// of partition X cites memory Y" intelligently.
    SummaryInput {
        /// Subject of the citing summary.
        from_subject: SummarySubject,
    },
    /// `child → parent` partition edge (or `memory → partition`).
    ParentPartition,
    /// `parent → child` partition edge (or `partition → memory`).
    PartitionContains,
}

impl EdgeKind {
    /// Discriminant tag for [`TraverseOpts::include_kinds`] filtering.
    #[must_use]
    pub fn tag(&self) -> EdgeKindTag {
        match self {
            EdgeKind::Link => EdgeKindTag::Link,
            EdgeKind::SummaryInput { .. } => EdgeKindTag::SummaryInput,
            EdgeKind::ParentPartition => EdgeKindTag::ParentPartition,
            EdgeKind::PartitionContains => EdgeKindTag::PartitionContains,
        }
    }
}

/// Discriminant-only enum used as a `BTreeSet` key in [`TraverseOpts`].
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum EdgeKindTag {
    /// `memory ↔ memory` link.
    Link,
    /// `summary → input` citation.
    SummaryInput,
    /// `child → parent` partition.
    ParentPartition,
    /// `parent → child` partition.
    PartitionContains,
}

impl EdgeKindTag {
    /// Every variant. Used as the default for [`TraverseOpts::include_kinds`].
    #[must_use]
    pub fn all() -> BTreeSet<EdgeKindTag> {
        let mut s = BTreeSet::new();
        s.insert(EdgeKindTag::Link);
        s.insert(EdgeKindTag::SummaryInput);
        s.insert(EdgeKindTag::ParentPartition);
        s.insert(EdgeKindTag::PartitionContains);
        s
    }
}

/// One edge.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct GraphEdge {
    /// Source node.
    pub from: NodeRef,
    /// Destination node.
    pub to: NodeRef,
    /// Edge kind.
    pub kind: EdgeKind,
}

/// One node — currently a thin wrapper around `NodeRef` so future fields
/// (cached labels, attributes) can land additively.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct GraphNode {
    /// The node's ref.
    pub r#ref: NodeRef,
}

/// A traversal result. Nodes are de-duplicated; edges may include duplicates
/// when both directions of a link were traversed.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct Graph {
    /// Nodes in the graph.
    pub nodes: Vec<GraphNode>,
    /// Edges in the graph.
    pub edges: Vec<GraphEdge>,
}

/// Options for [`crate::Memory::traverse`].
#[non_exhaustive]
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(default)]
pub struct TraverseOpts {
    /// Which edge kinds to follow. Defaults to every kind via
    /// [`EdgeKindTag::all`].
    pub include_kinds: BTreeSet<EdgeKindTag>,
    /// Hard cap on the total number of nodes in the result. Default 256.
    pub max_nodes: u32,
    /// When false, parent / contains edges are never followed even if
    /// `include_kinds` lists them. Convenience for "memory + link only"
    /// queries that don't want the partition tree noise. Default `true`.
    pub include_partition_edges: bool,
}

impl Default for TraverseOpts {
    fn default() -> Self {
        Self {
            include_kinds: EdgeKindTag::all(),
            max_nodes: 256,
            include_partition_edges: true,
        }
    }
}

impl TraverseOpts {
    /// Whether the supplied edge tag is permitted by the current opts.
    #[must_use]
    pub fn permits(&self, tag: EdgeKindTag) -> bool {
        if !self.include_kinds.contains(&tag) {
            return false;
        }
        if !self.include_partition_edges
            && matches!(
                tag,
                EdgeKindTag::ParentPartition | EdgeKindTag::PartitionContains
            )
        {
            return false;
        }
        true
    }
}

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

    #[test]
    fn default_opts_permit_everything() {
        let o = TraverseOpts::default();
        assert!(o.permits(EdgeKindTag::Link));
        assert!(o.permits(EdgeKindTag::SummaryInput));
        assert!(o.permits(EdgeKindTag::ParentPartition));
        assert!(o.permits(EdgeKindTag::PartitionContains));
    }

    #[test]
    fn include_partition_edges_off_blocks_partition_tags() {
        let o = TraverseOpts {
            include_partition_edges: false,
            ..TraverseOpts::default()
        };
        assert!(o.permits(EdgeKindTag::Link));
        assert!(!o.permits(EdgeKindTag::ParentPartition));
        assert!(!o.permits(EdgeKindTag::PartitionContains));
    }

    #[test]
    fn graph_serdes() {
        let g = Graph::default();
        let s = serde_json::to_string(&g).unwrap();
        let _: Graph = serde_json::from_str(&s).unwrap();
    }
}