Skip to main content

codemem_core/
traits.rs

1use serde::{Deserialize, Serialize};
2use std::collections::HashMap;
3
4use crate::{CodememError, Edge, GraphNode, MemoryNode, Session};
5
6// ── Traits ──────────────────────────────────────────────────────────────────
7
8/// Vector backend trait for HNSW index operations.
9pub trait VectorBackend: Send + Sync {
10    /// Insert a vector with associated ID.
11    fn insert(&mut self, id: &str, embedding: &[f32]) -> Result<(), CodememError>;
12
13    /// Batch insert vectors.
14    fn insert_batch(&mut self, items: &[(String, Vec<f32>)]) -> Result<(), CodememError>;
15
16    /// Search for k nearest neighbors. Returns (id, distance) pairs.
17    fn search(&self, query: &[f32], k: usize) -> Result<Vec<(String, f32)>, CodememError>;
18
19    /// Remove a vector by ID.
20    fn remove(&mut self, id: &str) -> Result<bool, CodememError>;
21
22    /// Save the index to disk.
23    fn save(&self, path: &std::path::Path) -> Result<(), CodememError>;
24
25    /// Load the index from disk.
26    fn load(&mut self, path: &std::path::Path) -> Result<(), CodememError>;
27
28    /// Get index statistics.
29    fn stats(&self) -> VectorStats;
30}
31
32/// Statistics about the vector index.
33#[derive(Debug, Clone, Default, Serialize, Deserialize)]
34pub struct VectorStats {
35    pub count: usize,
36    pub dimensions: usize,
37    pub metric: String,
38    pub memory_bytes: usize,
39}
40
41/// Graph backend trait for graph operations.
42pub trait GraphBackend: Send + Sync {
43    /// Add a node to the graph.
44    fn add_node(&mut self, node: GraphNode) -> Result<(), CodememError>;
45
46    /// Get a node by ID.
47    fn get_node(&self, id: &str) -> Result<Option<GraphNode>, CodememError>;
48
49    /// Remove a node by ID.
50    fn remove_node(&mut self, id: &str) -> Result<bool, CodememError>;
51
52    /// Add an edge between two nodes.
53    fn add_edge(&mut self, edge: Edge) -> Result<(), CodememError>;
54
55    /// Get edges from a node.
56    fn get_edges(&self, node_id: &str) -> Result<Vec<Edge>, CodememError>;
57
58    /// Remove an edge by ID.
59    fn remove_edge(&mut self, id: &str) -> Result<bool, CodememError>;
60
61    /// BFS traversal from a start node up to max_depth.
62    fn bfs(&self, start_id: &str, max_depth: usize) -> Result<Vec<GraphNode>, CodememError>;
63
64    /// DFS traversal from a start node up to max_depth.
65    fn dfs(&self, start_id: &str, max_depth: usize) -> Result<Vec<GraphNode>, CodememError>;
66
67    /// Shortest path between two nodes.
68    fn shortest_path(&self, from: &str, to: &str) -> Result<Vec<String>, CodememError>;
69
70    /// Get graph statistics.
71    fn stats(&self) -> GraphStats;
72}
73
74/// Statistics about the graph.
75#[derive(Debug, Clone, Default, Serialize, Deserialize)]
76pub struct GraphStats {
77    pub node_count: usize,
78    pub edge_count: usize,
79    pub node_kind_counts: HashMap<String, usize>,
80    pub relationship_type_counts: HashMap<String, usize>,
81}
82
83// ── Storage Stats & Consolidation Types ─────────────────────────────────
84
85/// Database statistics.
86#[derive(Debug, Clone, Serialize, Deserialize)]
87pub struct StorageStats {
88    pub memory_count: usize,
89    pub embedding_count: usize,
90    pub node_count: usize,
91    pub edge_count: usize,
92}
93
94/// A single consolidation log entry.
95#[derive(Debug, Clone)]
96pub struct ConsolidationLogEntry {
97    pub cycle_type: String,
98    pub run_at: i64,
99    pub affected_count: usize,
100}
101
102// ── Storage Backend Trait ───────────────────────────────────────────────
103
104/// Pluggable storage backend trait for all persistence operations.
105///
106/// This trait unifies memory CRUD, embedding persistence, graph node/edge
107/// storage, sessions, consolidation, and pattern detection behind a single
108/// interface. Implementations include SQLite (default) and can be extended
109/// for SurrealDB, FalkorDB, or other backends.
110pub trait StorageBackend: Send + Sync {
111    // ── Memory CRUD ─────────────────────────────────────────────────
112
113    /// Insert a new memory. Returns Err(Duplicate) if content hash already exists.
114    fn insert_memory(&self, memory: &MemoryNode) -> Result<(), CodememError>;
115
116    /// Get a memory by ID. Updates access_count and last_accessed_at.
117    fn get_memory(&self, id: &str) -> Result<Option<MemoryNode>, CodememError>;
118
119    /// Get multiple memories by IDs in a single batch operation.
120    fn get_memories_batch(&self, ids: &[&str]) -> Result<Vec<MemoryNode>, CodememError>;
121
122    /// Update a memory's content and optionally its importance. Re-computes content hash.
123    fn update_memory(
124        &self,
125        id: &str,
126        content: &str,
127        importance: Option<f64>,
128    ) -> Result<(), CodememError>;
129
130    /// Delete a memory by ID. Returns true if a row was deleted.
131    fn delete_memory(&self, id: &str) -> Result<bool, CodememError>;
132
133    /// List all memory IDs, ordered by created_at descending.
134    fn list_memory_ids(&self) -> Result<Vec<String>, CodememError>;
135
136    /// List memory IDs scoped to a specific namespace.
137    fn list_memory_ids_for_namespace(&self, namespace: &str) -> Result<Vec<String>, CodememError>;
138
139    /// List all distinct namespaces.
140    fn list_namespaces(&self) -> Result<Vec<String>, CodememError>;
141
142    /// Get total memory count.
143    fn memory_count(&self) -> Result<usize, CodememError>;
144
145    // ── Embedding Persistence ───────────────────────────────────────
146
147    /// Store an embedding vector for a memory.
148    fn store_embedding(&self, memory_id: &str, embedding: &[f32]) -> Result<(), CodememError>;
149
150    /// Get an embedding by memory ID.
151    fn get_embedding(&self, memory_id: &str) -> Result<Option<Vec<f32>>, CodememError>;
152
153    /// Delete an embedding by memory ID. Returns true if a row was deleted.
154    fn delete_embedding(&self, memory_id: &str) -> Result<bool, CodememError>;
155
156    /// List all stored embeddings as (memory_id, embedding_vector) pairs.
157    fn list_all_embeddings(&self) -> Result<Vec<(String, Vec<f32>)>, CodememError>;
158
159    // ── Graph Node/Edge Persistence ─────────────────────────────────
160
161    /// Insert or replace a graph node.
162    fn insert_graph_node(&self, node: &GraphNode) -> Result<(), CodememError>;
163
164    /// Get a graph node by ID.
165    fn get_graph_node(&self, id: &str) -> Result<Option<GraphNode>, CodememError>;
166
167    /// Delete a graph node by ID. Returns true if a row was deleted.
168    fn delete_graph_node(&self, id: &str) -> Result<bool, CodememError>;
169
170    /// Get all graph nodes.
171    fn all_graph_nodes(&self) -> Result<Vec<GraphNode>, CodememError>;
172
173    /// Insert or replace a graph edge.
174    fn insert_graph_edge(&self, edge: &Edge) -> Result<(), CodememError>;
175
176    /// Get all edges from or to a node.
177    fn get_edges_for_node(&self, node_id: &str) -> Result<Vec<Edge>, CodememError>;
178
179    /// Get all graph edges.
180    fn all_graph_edges(&self) -> Result<Vec<Edge>, CodememError>;
181
182    /// Delete all graph edges connected to a node. Returns count deleted.
183    fn delete_graph_edges_for_node(&self, node_id: &str) -> Result<usize, CodememError>;
184
185    // ── Sessions ────────────────────────────────────────────────────
186
187    /// Start a new session.
188    fn start_session(&self, id: &str, namespace: Option<&str>) -> Result<(), CodememError>;
189
190    /// End a session with optional summary.
191    fn end_session(&self, id: &str, summary: Option<&str>) -> Result<(), CodememError>;
192
193    /// List sessions, optionally filtered by namespace, up to limit.
194    fn list_sessions(
195        &self,
196        namespace: Option<&str>,
197        limit: usize,
198    ) -> Result<Vec<Session>, CodememError>;
199
200    // ── Consolidation ───────────────────────────────────────────────
201
202    /// Record a consolidation run.
203    fn insert_consolidation_log(
204        &self,
205        cycle_type: &str,
206        affected_count: usize,
207    ) -> Result<(), CodememError>;
208
209    /// Get the last consolidation run for each cycle type.
210    fn last_consolidation_runs(&self) -> Result<Vec<ConsolidationLogEntry>, CodememError>;
211
212    // ── Pattern Detection Queries ───────────────────────────────────
213
214    /// Find repeated search patterns. Returns (pattern, count, memory_ids).
215    fn get_repeated_searches(
216        &self,
217        min_count: usize,
218        namespace: Option<&str>,
219    ) -> Result<Vec<(String, usize, Vec<String>)>, CodememError>;
220
221    /// Find file hotspots. Returns (file_path, count, memory_ids).
222    fn get_file_hotspots(
223        &self,
224        min_count: usize,
225        namespace: Option<&str>,
226    ) -> Result<Vec<(String, usize, Vec<String>)>, CodememError>;
227
228    /// Get tool usage statistics. Returns (tool_name, count) pairs.
229    fn get_tool_usage_stats(
230        &self,
231        namespace: Option<&str>,
232    ) -> Result<Vec<(String, usize)>, CodememError>;
233
234    /// Find decision chains. Returns (file_path, count, memory_ids).
235    fn get_decision_chains(
236        &self,
237        min_count: usize,
238        namespace: Option<&str>,
239    ) -> Result<Vec<(String, usize, Vec<String>)>, CodememError>;
240
241    // ── Bulk Operations ─────────────────────────────────────────────
242
243    /// Decay importance of stale memories older than threshold_ts by decay_factor.
244    /// Returns count of affected memories.
245    fn decay_stale_memories(
246        &self,
247        threshold_ts: i64,
248        decay_factor: f64,
249    ) -> Result<usize, CodememError>;
250
251    /// List memories for creative consolidation: (id, memory_type, tags).
252    fn list_memories_for_creative(
253        &self,
254    ) -> Result<Vec<(String, String, Vec<String>)>, CodememError>;
255
256    /// Find near-duplicate memories by content hash prefix similarity.
257    /// Returns (id1, id2, similarity) pairs.
258    fn find_cluster_duplicates(&self) -> Result<Vec<(String, String, f64)>, CodememError>;
259
260    /// Find memories eligible for forgetting (low importance).
261    /// Returns list of memory IDs.
262    fn find_forgettable(&self, importance_threshold: f64) -> Result<Vec<String>, CodememError>;
263
264    // ── Batch Operations ────────────────────────────────────────────
265
266    /// Insert multiple memories in a single batch. Default impl calls insert_memory in a loop.
267    fn insert_memories_batch(&self, memories: &[MemoryNode]) -> Result<(), CodememError> {
268        for memory in memories {
269            self.insert_memory(memory)?;
270        }
271        Ok(())
272    }
273
274    /// Store multiple embeddings in a single batch. Default impl calls store_embedding in a loop.
275    fn store_embeddings_batch(&self, items: &[(&str, &[f32])]) -> Result<(), CodememError> {
276        for (id, embedding) in items {
277            self.store_embedding(id, embedding)?;
278        }
279        Ok(())
280    }
281
282    /// Insert multiple graph nodes in a single batch. Default impl calls insert_graph_node in a loop.
283    fn insert_graph_nodes_batch(&self, nodes: &[GraphNode]) -> Result<(), CodememError> {
284        for node in nodes {
285            self.insert_graph_node(node)?;
286        }
287        Ok(())
288    }
289
290    /// Insert multiple graph edges in a single batch. Default impl calls insert_graph_edge in a loop.
291    fn insert_graph_edges_batch(&self, edges: &[Edge]) -> Result<(), CodememError> {
292        for edge in edges {
293            self.insert_graph_edge(edge)?;
294        }
295        Ok(())
296    }
297
298    // ── Query Helpers ───────────────────────────────────────────────
299
300    /// Find memories that have no embeddings yet. Returns (id, content) pairs.
301    fn find_unembedded_memories(&self) -> Result<Vec<(String, String)>, CodememError>;
302
303    /// Search graph nodes by label (case-insensitive LIKE). Returns matching nodes
304    /// sorted by centrality descending, limited to `limit` results.
305    fn search_graph_nodes(
306        &self,
307        query: &str,
308        namespace: Option<&str>,
309        limit: usize,
310    ) -> Result<Vec<GraphNode>, CodememError>;
311
312    /// List memories with optional namespace and memory_type filters.
313    fn list_memories_filtered(
314        &self,
315        namespace: Option<&str>,
316        memory_type: Option<&str>,
317    ) -> Result<Vec<MemoryNode>, CodememError>;
318
319    /// Get edges filtered by namespace (edges where both src and dst nodes have the given namespace).
320    fn graph_edges_for_namespace(&self, namespace: &str) -> Result<Vec<Edge>, CodememError>;
321
322    // ── File Hash Tracking ──────────────────────────────────────────
323
324    /// Load all file hashes for incremental indexing. Returns path -> hash map.
325    fn load_file_hashes(&self) -> Result<HashMap<String, String>, CodememError>;
326
327    /// Save file hashes for incremental indexing.
328    fn save_file_hashes(&self, hashes: &HashMap<String, String>) -> Result<(), CodememError>;
329
330    // ── Stats ───────────────────────────────────────────────────────
331
332    /// Get database statistics.
333    fn stats(&self) -> Result<StorageStats, CodememError>;
334}