agentic_memory/graph/
builder.rs

1//! Fluent API for building MemoryGraph instances.
2
3use crate::types::{
4    AmemResult, CognitiveEvent, CognitiveEventBuilder, Edge, EdgeType, EventType, DEFAULT_DIMENSION,
5};
6
7use super::MemoryGraph;
8
9/// Fluent builder for constructing a MemoryGraph.
10pub struct GraphBuilder {
11    dimension: usize,
12    nodes: Vec<CognitiveEvent>,
13    edges: Vec<Edge>,
14    next_id: u64,
15}
16
17impl GraphBuilder {
18    /// Create a new builder with the default dimension.
19    pub fn new() -> Self {
20        Self {
21            dimension: DEFAULT_DIMENSION,
22            nodes: Vec::new(),
23            edges: Vec::new(),
24            next_id: 0,
25        }
26    }
27
28    /// Create a new builder with a specific dimension.
29    pub fn with_dimension(dim: usize) -> Self {
30        Self {
31            dimension: dim,
32            nodes: Vec::new(),
33            edges: Vec::new(),
34            next_id: 0,
35        }
36    }
37
38    fn add_event(
39        &mut self,
40        event_type: EventType,
41        content: &str,
42        session_id: u32,
43        confidence: f32,
44    ) -> u64 {
45        let id = self.next_id;
46        self.next_id += 1;
47        let event = CognitiveEventBuilder::new(event_type, content)
48            .session_id(session_id)
49            .confidence(confidence)
50            .feature_vec(vec![0.0; self.dimension])
51            .build();
52        let mut event = event;
53        event.id = id;
54        self.nodes.push(event);
55        id
56    }
57
58    /// Add a fact.
59    pub fn add_fact(&mut self, content: &str, session_id: u32, confidence: f32) -> u64 {
60        self.add_event(EventType::Fact, content, session_id, confidence)
61    }
62
63    /// Add a decision.
64    pub fn add_decision(&mut self, content: &str, session_id: u32, confidence: f32) -> u64 {
65        self.add_event(EventType::Decision, content, session_id, confidence)
66    }
67
68    /// Add an inference.
69    pub fn add_inference(&mut self, content: &str, session_id: u32, confidence: f32) -> u64 {
70        self.add_event(EventType::Inference, content, session_id, confidence)
71    }
72
73    /// Add a correction (automatically creates SUPERSEDES edge to old_node_id).
74    pub fn add_correction(&mut self, content: &str, session_id: u32, old_node_id: u64) -> u64 {
75        let id = self.add_event(EventType::Correction, content, session_id, 1.0);
76        self.edges
77            .push(Edge::new(id, old_node_id, EdgeType::Supersedes, 1.0));
78        // Reduce old node's confidence
79        if let Some(old_node) = self.nodes.iter_mut().find(|n| n.id == old_node_id) {
80            old_node.confidence = 0.0;
81        }
82        id
83    }
84
85    /// Add a skill.
86    pub fn add_skill(&mut self, content: &str, session_id: u32, confidence: f32) -> u64 {
87        self.add_event(EventType::Skill, content, session_id, confidence)
88    }
89
90    /// Add an episode summary (automatically creates PART_OF edges).
91    pub fn add_episode(&mut self, content: &str, session_id: u32, member_node_ids: &[u64]) -> u64 {
92        let id = self.add_event(EventType::Episode, content, session_id, 1.0);
93        for &member_id in member_node_ids {
94            self.edges
95                .push(Edge::new(member_id, id, EdgeType::PartOf, 1.0));
96        }
97        id
98    }
99
100    /// Add an edge between two nodes.
101    pub fn link(
102        &mut self,
103        source_id: u64,
104        target_id: u64,
105        edge_type: EdgeType,
106        weight: f32,
107    ) -> &mut Self {
108        self.edges
109            .push(Edge::new(source_id, target_id, edge_type, weight));
110        self
111    }
112
113    /// Set feature vector for a node.
114    pub fn set_feature_vec(&mut self, node_id: u64, vec: Vec<f32>) -> &mut Self {
115        if let Some(node) = self.nodes.iter_mut().find(|n| n.id == node_id) {
116            node.feature_vec = vec;
117        }
118        self
119    }
120
121    /// Build the final MemoryGraph.
122    pub fn build(self) -> AmemResult<MemoryGraph> {
123        MemoryGraph::from_parts(self.nodes, self.edges, self.dimension)
124    }
125}
126
127impl Default for GraphBuilder {
128    fn default() -> Self {
129        Self::new()
130    }
131}
agentic_memory/graph/builder.rs

agentic_memory/graph/
builder.rs
