Skip to main content

nusy_arrow_core/
graph_factory.rs

1//! Graph store factory — select and create the right store backend.
2//!
3//! Replaces Python `brain/knowledge/graph_factory.py`. The Python factory
4//! selected between GPU (cuDF) and CPU (Polars) backends. In V14 Rust,
5//! all backends use Arrow. The factory selects the right *level* of store
6//! based on use case and persistence requirements.
7//!
8//! # Backends
9//!
10//! | Backend | Store Type | Persistence | Use Case |
11//! |---------|-----------|-------------|----------|
12//! | `InMemory` | `ArrowGraphStore` | None | Tests, short-lived queries |
13//! | `KnowledgeGraph` | `KgStore` | None | Prefix-aware KG operations |
14//! | `Simple` | `SimpleTripleStore` | None | Default namespace + Y-layer |
15//!
16//! For persistent (hot/cold + Parquet) backends, use `nusy-dual-store::DualStore`
17//! directly — it has its own `DualStoreConfig` builder.
18//!
19//! # Hardware Detection
20//!
21//! `detect_backends()` reports available compute capabilities. Currently all
22//! backends are CPU (Arrow). GPU acceleration (CUDA via Candle) is tracked
23//! separately in V6 and does not affect graph store selection.
24
25use crate::kg_store::KgStore;
26use crate::namespace::Namespace;
27use crate::store::ArrowGraphStore;
28use crate::triple_store::SimpleTripleStore;
29use crate::y_layer::YLayer;
30
31// ── Backend selection ───────────────────────────────────────────────────
32
33/// Available graph store backends.
34#[derive(Debug, Clone, Copy, PartialEq, Eq)]
35pub enum GraphBackend {
36    /// Raw Arrow graph store — no defaults, no extras.
37    /// Best for: tests, manual namespace management, maximum control.
38    InMemory,
39
40    /// Simple triple store with default namespace + Y-layer.
41    /// Best for: single-namespace work, quick prototyping.
42    Simple,
43
44    /// Full knowledge graph with prefix management and keyword search.
45    /// Best for: being knowledge, ontology work, gap tracking.
46    KnowledgeGraph,
47}
48
49/// Hardware capabilities detected on this machine.
50#[derive(Debug, Clone)]
51pub struct HardwareCapabilities {
52    /// Whether CUDA GPU is available (DGX).
53    pub cuda_available: bool,
54    /// Whether MPS (Metal) is available (M4/M5 Macs).
55    pub mps_available: bool,
56    /// Number of CPU cores.
57    pub cpu_cores: usize,
58    /// Approximate available memory in bytes.
59    pub memory_bytes: u64,
60}
61
62/// Detect available hardware capabilities.
63///
64/// Currently reports CPU-only for graph stores. GPU availability is
65/// informational — graph stores don't use GPU (training does, via V6).
66pub fn detect_hardware() -> HardwareCapabilities {
67    let cpu_cores = std::thread::available_parallelism()
68        .map(|p| p.get())
69        .unwrap_or(1);
70
71    // CUDA detection: check for nvidia-smi
72    let cuda_available = std::process::Command::new("nvidia-smi")
73        .arg("--query-gpu=name")
74        .arg("--format=csv,noheader")
75        .output()
76        .map(|o| o.status.success())
77        .unwrap_or(false);
78
79    // MPS detection: macOS with Apple Silicon
80    let mps_available = cfg!(target_os = "macos") && cfg!(target_arch = "aarch64");
81
82    HardwareCapabilities {
83        cuda_available,
84        mps_available,
85        cpu_cores,
86        memory_bytes: 0, // Not easily portable; callers can override
87    }
88}
89
90/// List available graph backends on this machine.
91///
92/// All three backends are always available (they're CPU Arrow).
93/// This function exists for parity with Python's `get_available_backends()`.
94pub fn available_backends() -> Vec<GraphBackend> {
95    vec![
96        GraphBackend::InMemory,
97        GraphBackend::Simple,
98        GraphBackend::KnowledgeGraph,
99    ]
100}
101
102/// Select the recommended backend based on use case.
103///
104/// This is the main factory entry point. For most being work,
105/// `KnowledgeGraph` is the right choice.
106pub fn recommended_backend() -> GraphBackend {
107    GraphBackend::KnowledgeGraph
108}
109
110// ── Factory configuration ───────────────────────────────────────────────
111
112/// Configuration for graph store creation.
113#[derive(Debug, Clone)]
114pub struct GraphStoreConfig {
115    /// Which backend to use.
116    pub backend: GraphBackend,
117    /// Default namespace (for Simple and KnowledgeGraph backends).
118    pub default_namespace: Namespace,
119    /// Default Y-layer (for Simple and KnowledgeGraph backends).
120    pub default_y_layer: YLayer,
121}
122
123impl Default for GraphStoreConfig {
124    fn default() -> Self {
125        Self {
126            backend: GraphBackend::KnowledgeGraph,
127            default_namespace: Namespace::World,
128            default_y_layer: YLayer::Semantic,
129        }
130    }
131}
132
133impl GraphStoreConfig {
134    /// Create config for a specific backend.
135    pub fn new(backend: GraphBackend) -> Self {
136        Self {
137            backend,
138            ..Default::default()
139        }
140    }
141
142    /// Set the default namespace.
143    pub fn with_namespace(mut self, ns: Namespace) -> Self {
144        self.default_namespace = ns;
145        self
146    }
147
148    /// Set the default Y-layer.
149    pub fn with_y_layer(mut self, layer: YLayer) -> Self {
150        self.default_y_layer = layer;
151        self
152    }
153}
154
155// ── Factory functions ───────────────────────────────────────────────────
156
157/// Created graph store — enum dispatch over backend types.
158///
159/// This avoids trait objects while still providing a unified return type.
160/// Callers match on the variant to get the concrete store.
161pub enum CreatedStore {
162    InMemory(ArrowGraphStore),
163    Simple(SimpleTripleStore),
164    KnowledgeGraph(KgStore),
165}
166
167impl CreatedStore {
168    /// Get the number of triples in the store.
169    pub fn len(&self) -> usize {
170        match self {
171            Self::InMemory(s) => s.len(),
172            Self::Simple(s) => s.len(),
173            Self::KnowledgeGraph(s) => s.len(),
174        }
175    }
176
177    /// Whether the store is empty.
178    pub fn is_empty(&self) -> bool {
179        self.len() == 0
180    }
181
182    /// Unwrap as ArrowGraphStore (panics if wrong variant).
183    pub fn into_arrow(self) -> ArrowGraphStore {
184        match self {
185            Self::InMemory(s) => s,
186            _ => panic!("expected InMemory variant"),
187        }
188    }
189
190    /// Unwrap as SimpleTripleStore (panics if wrong variant).
191    pub fn into_simple(self) -> SimpleTripleStore {
192        match self {
193            Self::Simple(s) => s,
194            _ => panic!("expected Simple variant"),
195        }
196    }
197
198    /// Unwrap as KgStore (panics if wrong variant).
199    pub fn into_kg(self) -> KgStore {
200        match self {
201            Self::KnowledgeGraph(s) => s,
202            _ => panic!("expected KnowledgeGraph variant"),
203        }
204    }
205
206    /// Try to unwrap as ArrowGraphStore.
207    pub fn try_into_arrow(self) -> Option<ArrowGraphStore> {
208        match self {
209            Self::InMemory(s) => Some(s),
210            _ => None,
211        }
212    }
213
214    /// Try to unwrap as SimpleTripleStore.
215    pub fn try_into_simple(self) -> Option<SimpleTripleStore> {
216        match self {
217            Self::Simple(s) => Some(s),
218            _ => None,
219        }
220    }
221
222    /// Try to unwrap as KgStore.
223    pub fn try_into_kg(self) -> Option<KgStore> {
224        match self {
225            Self::KnowledgeGraph(s) => Some(s),
226            _ => None,
227        }
228    }
229}
230
231/// Create a graph store based on configuration.
232///
233/// This is the primary factory function. Equivalent to Python's
234/// `create_graph_store(backend="auto")`.
235pub fn create_graph_store(config: &GraphStoreConfig) -> CreatedStore {
236    match config.backend {
237        GraphBackend::InMemory => CreatedStore::InMemory(ArrowGraphStore::new()),
238        GraphBackend::Simple => CreatedStore::Simple(SimpleTripleStore::with_defaults(
239            config.default_namespace,
240            config.default_y_layer,
241        )),
242        GraphBackend::KnowledgeGraph => CreatedStore::KnowledgeGraph(KgStore::with_defaults(
243            config.default_namespace,
244            config.default_y_layer,
245        )),
246    }
247}
248
249/// Create a graph store with default configuration (KnowledgeGraph backend).
250pub fn create_default_store() -> KgStore {
251    KgStore::new()
252}
253
254// ── Tests ───────────────────────────────────────────────────────────────
255
256#[cfg(test)]
257mod tests {
258    use super::*;
259
260    #[test]
261    fn test_detect_hardware() {
262        let hw = detect_hardware();
263        assert!(hw.cpu_cores >= 1);
264        // MPS should be true on Apple Silicon Macs
265        if cfg!(target_os = "macos") && cfg!(target_arch = "aarch64") {
266            assert!(hw.mps_available);
267        }
268    }
269
270    #[test]
271    fn test_available_backends() {
272        let backends = available_backends();
273        assert_eq!(backends.len(), 3);
274        assert!(backends.contains(&GraphBackend::InMemory));
275        assert!(backends.contains(&GraphBackend::Simple));
276        assert!(backends.contains(&GraphBackend::KnowledgeGraph));
277    }
278
279    #[test]
280    fn test_recommended_backend() {
281        assert_eq!(recommended_backend(), GraphBackend::KnowledgeGraph);
282    }
283
284    #[test]
285    fn test_create_in_memory() {
286        let config = GraphStoreConfig::new(GraphBackend::InMemory);
287        let store = create_graph_store(&config);
288        assert!(store.is_empty());
289        let arrow = store.into_arrow();
290        assert_eq!(arrow.len(), 0);
291    }
292
293    #[test]
294    fn test_create_simple() {
295        let config = GraphStoreConfig::new(GraphBackend::Simple)
296            .with_namespace(Namespace::Research)
297            .with_y_layer(YLayer::Reasoning);
298        let store = create_graph_store(&config);
299        assert!(store.is_empty());
300        let simple = store.into_simple();
301        assert_eq!(simple.len(), 0);
302    }
303
304    #[test]
305    fn test_create_knowledge_graph() {
306        let config = GraphStoreConfig::new(GraphBackend::KnowledgeGraph);
307        let store = create_graph_store(&config);
308        assert!(store.is_empty());
309        let kg = store.into_kg();
310        // KgStore comes with default prefixes
311        assert!(!kg.prefixes().is_empty());
312    }
313
314    #[test]
315    fn test_create_default_store() {
316        let store = create_default_store();
317        assert!(store.is_empty());
318        assert!(!store.prefixes().is_empty());
319    }
320
321    #[test]
322    fn test_default_config() {
323        let config = GraphStoreConfig::default();
324        assert_eq!(config.backend, GraphBackend::KnowledgeGraph);
325        assert_eq!(config.default_namespace, Namespace::World);
326        assert_eq!(config.default_y_layer, YLayer::Semantic);
327    }
328
329    #[test]
330    fn test_try_into_wrong_variant() {
331        let store = create_graph_store(&GraphStoreConfig::new(GraphBackend::InMemory));
332        assert!(store.try_into_kg().is_none());
333
334        let store = create_graph_store(&GraphStoreConfig::new(GraphBackend::KnowledgeGraph));
335        assert!(store.try_into_arrow().is_none());
336    }
337
338    #[test]
339    fn test_created_store_len() {
340        let store = create_graph_store(&GraphStoreConfig::new(GraphBackend::InMemory));
341        assert_eq!(store.len(), 0);
342        assert!(store.is_empty());
343    }
344
345    #[test]
346    fn test_graceful_fallback_no_gpu() {
347        // On Mini/M5 (no CUDA), factory should still work — all backends are CPU
348        let _hw = detect_hardware();
349        // Regardless of CUDA availability, all backends work
350        for backend in available_backends() {
351            let config = GraphStoreConfig::new(backend);
352            let store = create_graph_store(&config);
353            assert!(store.is_empty());
354        }
355    }
356}