agentic_memory/index/
session_index.rs

1//! Session index — maps each session_id to sorted node IDs in that session.
2
3use std::collections::HashMap;
4
5use crate::types::CognitiveEvent;
6
7/// Maps each session_id to a sorted list of node IDs in that session.
8pub struct SessionIndex {
9    index: HashMap<u32, Vec<u64>>,
10}
11
12impl SessionIndex {
13    /// Create a new, empty session index.
14    pub fn new() -> Self {
15        Self {
16            index: HashMap::new(),
17        }
18    }
19
20    /// Get all node IDs from a specific session.
21    pub fn get_session(&self, session_id: u32) -> &[u64] {
22        self.index
23            .get(&session_id)
24            .map(|v| v.as_slice())
25            .unwrap_or(&[])
26    }
27
28    /// Get all node IDs from multiple sessions, merged and sorted.
29    pub fn get_sessions(&self, session_ids: &[u32]) -> Vec<u64> {
30        let mut result: Vec<u64> = Vec::new();
31        for sid in session_ids {
32            if let Some(ids) = self.index.get(sid) {
33                result.extend_from_slice(ids);
34            }
35        }
36        result.sort_unstable();
37        result.dedup();
38        result
39    }
40
41    /// Get all known session IDs.
42    pub fn session_ids(&self) -> Vec<u32> {
43        let mut ids: Vec<u32> = self.index.keys().copied().collect();
44        ids.sort_unstable();
45        ids
46    }
47
48    /// Count sessions.
49    pub fn session_count(&self) -> usize {
50        self.index.len()
51    }
52
53    /// Count nodes in a session.
54    pub fn node_count(&self, session_id: u32) -> usize {
55        self.index.get(&session_id).map(|v| v.len()).unwrap_or(0)
56    }
57
58    /// Rebuild the entire index from a slice of nodes.
59    pub fn rebuild(&mut self, nodes: &[CognitiveEvent]) {
60        self.index.clear();
61        for node in nodes {
62            self.index.entry(node.session_id).or_default().push(node.id);
63        }
64        for list in self.index.values_mut() {
65            list.sort_unstable();
66        }
67    }
68
69    /// Incrementally add a new node.
70    pub fn add_node(&mut self, event: &CognitiveEvent) {
71        let list = self.index.entry(event.session_id).or_default();
72        let pos = list.binary_search(&event.id).unwrap_or_else(|p| p);
73        list.insert(pos, event.id);
74    }
75
76    /// Remove a node from the index.
77    pub fn remove_node(&mut self, id: u64, session_id: u32) {
78        if let Some(list) = self.index.get_mut(&session_id) {
79            if let Ok(pos) = list.binary_search(&id) {
80                list.remove(pos);
81            }
82            if list.is_empty() {
83                self.index.remove(&session_id);
84            }
85        }
86    }
87
88    /// Clear the index.
89    pub fn clear(&mut self) {
90        self.index.clear();
91    }
92
93    /// Number of total entries across all sessions.
94    pub fn len(&self) -> usize {
95        self.index.values().map(|v| v.len()).sum()
96    }
97
98    /// Whether the index is empty.
99    pub fn is_empty(&self) -> bool {
100        self.len() == 0
101    }
102
103    /// Get a reference to the underlying map (for serialization).
104    pub fn inner(&self) -> &HashMap<u32, Vec<u64>> {
105        &self.index
106    }
107}
108
109impl Default for SessionIndex {
110    fn default() -> Self {
111        Self::new()
112    }
113}
agentic_memory/index/session_index.rs

agentic_memory/index/
session_index.rs
