feat: add num_cpus dependency and improve code formatting

- Added `num_cpus` dependency to `Cargo.toml`.
- Refactored `http_server.rs`, `stdio_server.rs`, `config.rs`, `error.rs`, `lib.rs`, `rag_tools.rs`, and `tests.rs` for improved code formatting and consistency.
- Removed unnecessary blank lines and adjusted line breaks for better readability.
- Enhanced error messages in `CoreRagMcpServer` methods for clarity.

Author: jruokola

crates/codegraph-ai/Cargo.toml

@@ -25,3 +25,4 @@ futures = { workspace = true }
 similar = { workspace = true }
 arc-swap = { workspace = true }
 num_cpus = { workspace = true }
+prometheus = { workspace = true }

crates/codegraph-ai/src/lib.rs

@@ -1,4 +1,6 @@
+pub mod ml;
+pub mod optimization;
+pub mod rag;
 pub mod semantic;
 
 pub use semantic::search::*;
-

crates/codegraph-ai/src/ml/mod.rs

@@ -1,2 +1 @@
 pub mod pipeline;
-

crates/codegraph-ai/src/ml/pipeline.rs

@@ -25,7 +25,7 @@ use uuid::Uuid;
 
 use codegraph_core::{CodeNode, NodeId, Result};
 use codegraph_vector::ml as vml;
-use codegraph_vector::{EmbeddingGenerator};
+use codegraph_vector::EmbeddingGenerator;
 
 /// Versioned model metadata
 #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
@@ -46,10 +46,18 @@ pub struct ModelRegistry {
 
 impl ModelRegistry {
     pub fn new<P: Into<PathBuf>>(root: P) -> Self {
-        Self { root: root.into(), index: PLRwLock::new(HashMap::new()) }
+        Self {
+            root: root.into(),
+            index: PLRwLock::new(HashMap::new()),
+        }
     }
 
-    pub async fn register(&self, model_name: &str, version: &str, metrics: HashMap<String, f32>) -> Result<ModelVersionMeta> {
+    pub async fn register(
+        &self,
+        model_name: &str,
+        version: &str,
+        metrics: HashMap<String, f32>,
+    ) -> Result<ModelVersionMeta> {
         let dir = self.root.join(model_name).join(version);
         tokio::fs::create_dir_all(&dir).await.ok();
 
@@ -60,7 +68,9 @@ impl ModelRegistry {
             metrics,
             path: dir.clone(),
         };
-        self.index.write().entry(model_name.to_string())
+        self.index
+            .write()
+            .entry(model_name.to_string())
             .or_default()
             .insert(version.to_string(), meta.clone());
 
@@ -79,7 +89,11 @@ impl ModelRegistry {
     }
 
     pub fn get(&self, model_name: &str, version: &str) -> Option<ModelVersionMeta> {
-        self.index.read().get(model_name).and_then(|m| m.get(version)).cloned()
+        self.index
+            .read()
+            .get(model_name)
+            .and_then(|m| m.get(version))
+            .cloned()
     }
 }
 
@@ -91,11 +105,17 @@ pub struct HotSwapModel {
 
 impl HotSwapModel {
     pub fn new<S: Into<String>>(name: S, initial_version: S) -> Self {
-        Self { active_name: name.into(), active_version: ArcSwap::from_pointee(initial_version.into()) }
+        Self {
+            active_name: name.into(),
+            active_version: ArcSwap::from_pointee(initial_version.into()),
+        }
     }
 
     pub fn active(&self) -> (String, String) {
-        (self.active_name.clone(), (*self.active_version.load()).clone())
+        (
+            self.active_name.clone(),
+            (*self.active_version.load()).clone(),
+        )
     }
 
     pub fn swap_version<S: Into<String>>(&self, new_version: S) {
@@ -128,7 +148,12 @@ impl Default for AiPipelineBuilder {
                 hyperparameters: vml::TrainingHyperparameters::default(),
                 data_config: vml::DataConfig::default(),
                 validation_config: vml::ValidationConfig::default(),
-                output_config: vml::OutputConfig { model_path: "models".into(), save_checkpoints: true, checkpoint_frequency: 10, export_for_inference: true }
+                output_config: vml::OutputConfig {
+                    model_path: "models".into(),
+                    save_checkpoints: true,
+                    checkpoint_frequency: 10,
+                    export_for_inference: true,
+                },
             },
             inference: vml::InferenceConfig::default(),
             registry_root: PathBuf::from("models"),
@@ -139,14 +164,34 @@ impl Default for AiPipelineBuilder {
 }
 
 impl AiPipelineBuilder {
-    pub fn new() -> Self { Self::default() }
+    pub fn new() -> Self {
+        Self::default()
+    }
 
-    pub fn feature_config(mut self, cfg: vml::FeatureConfig) -> Self { self.feature = cfg; self }
-    pub fn training_config(mut self, cfg: vml::TrainingConfig) -> Self { self.training = cfg; self }
-    pub fn inference_config(mut self, cfg: vml::InferenceConfig) -> Self { self.inference = cfg; self }
-    pub fn registry_root<P: Into<PathBuf>>(mut self, root: P) -> Self { self.registry_root = root.into(); self }
-    pub fn model_name<S: Into<String>>(mut self, name: S) -> Self { self.model_name = name.into(); self }
-    pub fn initial_version<S: Into<String>>(mut self, v: S) -> Self { self.initial_version = v.into(); self }
+    pub fn feature_config(mut self, cfg: vml::FeatureConfig) -> Self {
+        self.feature = cfg;
+        self
+    }
+    pub fn training_config(mut self, cfg: vml::TrainingConfig) -> Self {
+        self.training = cfg;
+        self
+    }
+    pub fn inference_config(mut self, cfg: vml::InferenceConfig) -> Self {
+        self.inference = cfg;
+        self
+    }
+    pub fn registry_root<P: Into<PathBuf>>(mut self, root: P) -> Self {
+        self.registry_root = root.into();
+        self
+    }
+    pub fn model_name<S: Into<String>>(mut self, name: S) -> Self {
+        self.model_name = name.into();
+        self
+    }
+    pub fn initial_version<S: Into<String>>(mut self, v: S) -> Self {
+        self.initial_version = v.into();
+        self
+    }
 
     pub fn build(self) -> Result<AiPipeline> {
         let embedding_generator = Arc::new(EmbeddingGenerator::default());
@@ -161,51 +206,120 @@ impl AiPipelineBuilder {
         let registry = Arc::new(ModelRegistry::new(&self.registry_root));
         let active = Arc::new(HotSwapModel::new(&self.model_name, &self.initial_version));
 
-        Ok(AiPipeline { inner, registry, active })
+        Ok(AiPipeline {
+            inner,
+            registry,
+            active,
+        })
     }
 }
 
 impl AiPipeline {
-    pub fn builder() -> AiPipelineBuilder { AiPipelineBuilder::new() }
+    pub fn builder() -> AiPipelineBuilder {
+        AiPipelineBuilder::new()
+    }
 
     /// Initialize the inner pipeline
-    pub async fn initialize(&self) -> Result<()> { self.inner.initialize().await }
+    pub async fn initialize(&self) -> Result<()> {
+        self.inner.initialize().await
+    }
 
     /// Train and register a versioned model, then hot-swap as active if requested.
-    pub async fn train_and_deploy(&self, dataset: &str, nodes: &[CodeNode], targets: Vec<vml::TrainingTarget>, version: &str, set_active: bool) -> Result<vml::TrainingResults> {
-        let results = self.inner.train_model(dataset, nodes, targets, &self.active_model_name()).await?;
+    pub async fn train_and_deploy(
+        &self,
+        dataset: &str,
+        nodes: &[CodeNode],
+        targets: Vec<vml::TrainingTarget>,
+        version: &str,
+        set_active: bool,
+    ) -> Result<vml::TrainingResults> {
+        let results = self
+            .inner
+            .train_model(dataset, nodes, targets, &self.active_model_name())
+            .await?;
 
         // Register version
-        let meta = self.registry.register(
-            &self.active_model_name(),
-            version,
-            results.validation_metrics.clone(),
-        ).await?;
+        let meta = self
+            .registry
+            .register(
+                &self.active_model_name(),
+                version,
+                results.validation_metrics.clone(),
+            )
+            .await?;
 
         // Save model artifact
         let path = meta.path.join("model.json");
-        let _ = self.inner.save_model(&self.active_model_name(), &path).await;
+        let _ = self
+            .inner
+            .save_model(&self.active_model_name(), &path)
+            .await;
 
         // Hot swap
-        if set_active { self.active.swap_version(version.to_string()); }
+        if set_active {
+            self.active.swap_version(version.to_string());
+        }
         Ok(results)
     }
 
     /// Start an A/B test between two versions.
-    pub async fn start_ab_test(&self, experiment: &str, version_a: &str, version_b: &str, duration: Duration) -> Result<String> {
+    pub async fn start_ab_test(
+        &self,
+        experiment: &str,
+        version_a: &str,
+        version_b: &str,
+        duration: Duration,
+    ) -> Result<String> {
         // Ensure both versions exist
-        if self.registry.get(&self.active_model_name(), version_a).is_none() || self.registry.get(&self.active_model_name(), version_b).is_none() {
-            return Err(codegraph_core::CodeGraphError::Training("Model versions not found for A/B test".into()));
+        if self
+            .registry
+            .get(&self.active_model_name(), version_a)
+            .is_none()
+            || self
+                .registry
+                .get(&self.active_model_name(), version_b)
+                .is_none()
+        {
+            return Err(codegraph_core::CodeGraphError::Training(
+                "Model versions not found for A/B test".into(),
+            ));
         }
         let mut alloc = HashMap::new();
         alloc.insert("A".to_string(), 0.5);
         alloc.insert("B".to_string(), 0.5);
-        let traffic = vml::TrafficAllocation { allocations: alloc, strategy: vml::AllocationStrategy::WeightedRandom, sticky_sessions: true };
+        let traffic = vml::TrafficAllocation {
+            allocations: alloc,
+            strategy: vml::AllocationStrategy::WeightedRandom,
+            sticky_sessions: true,
+        };
         let stats = vml::StatisticalConfig::default();
-        let metrics = vec![vml::ExperimentMetric::Accuracy, vml::ExperimentMetric::Latency, vml::ExperimentMetric::Throughput];
-        let early = vml::EarlyStoppingConfig { enabled: true, check_interval: Duration::from_secs(60), min_samples: 500, futility_boundary: 0.01, efficacy_boundary: 0.01 };
-        let sample = vml::SampleSizeConfig { min_sample_size: 1000, max_sample_size: 100_000, early_stopping: early, calculation_method: vml::SampleSizeMethod::Sequential };
-        let cfg = vml::ABTestConfig { name: experiment.into(), description: "Model A/B comparison".into(), traffic_allocation: traffic, duration, statistical_config: stats, metrics, sample_size: sample };
+        let metrics = vec![
+            vml::ExperimentMetric::Accuracy,
+            vml::ExperimentMetric::Latency,
+            vml::ExperimentMetric::Throughput,
+        ];
+        let early = vml::EarlyStoppingConfig {
+            enabled: true,
+            check_interval: Duration::from_secs(60),
+            min_samples: 500,
+            futility_boundary: 0.01,
+            efficacy_boundary: 0.01,
+        };
+        let sample = vml::SampleSizeConfig {
+            min_sample_size: 1000,
+            max_sample_size: 100_000,
+            early_stopping: early,
+            calculation_method: vml::SampleSizeMethod::Sequential,
+        };
+        let cfg = vml::ABTestConfig {
+            name: experiment.into(),
+            description: "Model A/B comparison".into(),
+            traffic_allocation: traffic,
+            duration,
+            statistical_config: stats,
+            metrics,
+            sample_size: sample,
+        };
         let id = self.inner.start_ab_test(cfg).await?;
         Ok(id)
     }
@@ -217,30 +331,44 @@ impl AiPipeline {
     }
 
     /// High-throughput batch feature extraction (concurrent), returns features in input order.
-    pub async fn extract_features_batch_fast(&self, nodes: &[CodeNode]) -> Result<Vec<vml::CodeFeatures>> {
+    pub async fn extract_features_batch_fast(
+        &self,
+        nodes: &[CodeNode],
+    ) -> Result<Vec<vml::CodeFeatures>> {
         // Use the inner feature extractor via pipeline call; shard across tasks for concurrency
         let chunk = std::cmp::max(64, nodes.len() / std::cmp::max(1, num_cpus::get()));
         let mut tasks = Vec::new();
         for chunk_nodes in nodes.chunks(chunk) {
             let part = chunk_nodes.to_vec();
             let inner = self.inner.clone();
-            tasks.push(tokio::spawn(async move { inner.extract_features_batch(&part).await }));
+            tasks.push(tokio::spawn(async move {
+                inner.extract_features_batch(&part).await
+            }));
         }
         let mut out = Vec::with_capacity(nodes.len());
-        for t in tasks { out.extend(t.await.unwrap()?); }
+        for t in tasks {
+            out.extend(t.await.unwrap()?);
+        }
         Ok(out)
     }
 
     /// Active model name and version tuple
-    pub fn active(&self) -> (String, String) { self.active.active() }
-    pub fn active_model_name(&self) -> String { self.active.active().0 }
+    pub fn active(&self) -> (String, String) {
+        self.active.active()
+    }
+    pub fn active_model_name(&self) -> String {
+        self.active.active().0
+    }
 
     /// Zero-downtime deploy a new version: warm-up then hot-swap
     pub async fn deploy_version(&self, version: &str, warmup_samples: &[CodeNode]) -> Result<()> {
         // Load model artifact if needed (inner keeps in-memory models; ensure present)
         if let Some(meta) = self.registry.get(&self.active_model_name(), version) {
             let path = meta.path.join("model.json");
-            let _ = self.inner.load_model(&self.active_model_name(), &path).await; // best-effort
+            let _ = self
+                .inner
+                .load_model(&self.active_model_name(), &path)
+                .await; // best-effort
         }
 
         // Warm-up inference to prime caches and JIT paths
@@ -254,11 +382,17 @@ impl AiPipeline {
     }
 
     /// Expose inner metrics for monitoring SLA (latency, throughput, cache hit rate)
-    pub async fn metrics(&self) -> vml::InferenceMetrics { self.inner.get_inference_metrics().await }
+    pub async fn metrics(&self) -> vml::InferenceMetrics {
+        self.inner.get_inference_metrics().await
+    }
 
     /// Proxy helpers to inner pipeline for convenience
-    pub async fn save_config(&self, path: &Path) -> Result<()> { self.inner.save_config(path).await }
-    pub async fn load_config(&self, path: &Path) -> Result<()> { self.inner.load_config(path).await }
+    pub async fn save_config(&self, path: &Path) -> Result<()> {
+        self.inner.save_config(path).await
+    }
+    pub async fn load_config(&self, path: &Path) -> Result<()> {
+        self.inner.load_config(path).await
+    }
 }
 
 // Lightweight proxy methods on inner MLPipeline (implement Clone by arc-wrapping inside inner)
@@ -267,7 +401,8 @@ trait CloneablePipeline {
 }
 
 impl CloneablePipeline for vml::MLPipeline {
-    fn clone(&self) -> Self { // safe shallow rebuild via saved config and shared internals
+    fn clone(&self) -> Self {
+        // safe shallow rebuild via saved config and shared internals
         // Use builder + current config snapshot
         // Read-only operations in `build` path; acceptable for proxy clone
         let cfg = futures::executor::block_on(async { self.get_context().await.config.clone() });
@@ -291,7 +426,14 @@ mod tests {
         let p = AiPipeline::builder().build().unwrap();
         p.initialize().await.unwrap();
 
-        let node = CodeNode { id: "n1".into(), name: "foo".into(), language: Some(Language::Rust), node_type: Some(NodeType::Function), content: Some("fn foo() { 1 }".into()), children: None };
+        let node = CodeNode {
+            id: "n1".into(),
+            name: "foo".into(),
+            language: Some(Language::Rust),
+            node_type: Some(NodeType::Function),
+            content: Some("fn foo() { 1 }".into()),
+            children: None,
+        };
         let _ = p.infer(&node).await.unwrap();
     }
 }