prometheusnn 0.2.1__tar.gz → 0.2.2__tar.gz

{prometheusnn-0.2.1 → prometheusnn-0.2.2}/CHANGELOG.md

@@ -5,6 +5,18 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.2.2] - 2026-10-30
+
+### Security
+- **Novelty data is loaded without pickle.** `novelty_centroids.npz` was loaded
+  with `allow_pickle=True`, so a crafted model directory could run code in the
+  server when any tool loaded that model. In 0.2.0 and earlier this was
+  reachable through the `model_name` path traversal fixed in 0.2.1; in 0.2.1 it
+  still applied to a malicious model placed in the models directory. Both
+  novelty files now load with `allow_pickle=False`. Models saved by
+  PrometheusNN contain only numeric and string arrays and load unchanged; a
+  file holding pickled objects is refused with an error asking for a retrain.
+
 ## [0.2.1] - 2026-10-30
 
 ### Security

{prometheusnn-0.2.1 → prometheusnn-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: prometheusnn
-Version: 0.2.1
+Version: 0.2.2
 Summary: Frictionless hierarchical text classification — an MCP server that guides you from CSV to trained classifier
 Project-URL: Homepage, https://github.com/mfbaig35r/prometheusnn
 Project-URL: Repository, https://github.com/mfbaig35r/prometheusnn

{prometheusnn-0.2.1 → prometheusnn-0.2.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "prometheusnn"
-version = "0.2.1"
+version = "0.2.2"
 description = "Frictionless hierarchical text classification — an MCP server that guides you from CSV to trained classifier"
 readme = "README.md"
 requires-python = ">=3.11"

{prometheusnn-0.2.1 → prometheusnn-0.2.2}/src/prometheus/inference/novelty.py

@@ -116,17 +116,26 @@ class NoveltyDetector:
 
     def load(self, path: Path):
         path = Path(path)
-        data = np.load(path / "novelty_centroids.npz", allow_pickle=True)
-        class_names = data["class_names"]
-        mean_dists = data["mean_dists"]
-        std_dists = data["std_dists"]
-        p95_dists = data["p95_dists"]
-        counts = data["counts"]
+        # Never unpickle: a model directory is data, and a crafted .npz with object arrays would
+        # run code on load. save() writes only numeric and string arrays, which need no pickle.
+        data = np.load(path / "novelty_centroids.npz", allow_pickle=False)
+        try:
+            class_names = data["class_names"]
+            mean_dists = data["mean_dists"]
+            std_dists = data["std_dists"]
+            p95_dists = data["p95_dists"]
+            counts = data["counts"]
+            centroids = [data[f"centroid_{i}"] for i in range(len(class_names))]
+        except ValueError as exc:
+            raise ValueError(
+                f"{path / 'novelty_centroids.npz'} contains pickled objects; refusing to load it. "
+                "Retrain the model to regenerate it."
+            ) from exc
 
         self.centroids = {}
         for i, name in enumerate(class_names):
             self.centroids[str(name)] = ClassCentroid(
-                centroid=data[f"centroid_{i}"],
+                centroid=centroids[i],
                 mean_dist=float(mean_dists[i]),
                 std_dist=max(float(std_dists[i]), 0.1),
                 p95_dist=float(p95_dists[i]),
@@ -135,7 +144,7 @@ class NoveltyDetector:
 
         emb_path = path / "novelty_embeddings.npy"
         if emb_path.exists():
-            self.training_embeddings = np.load(emb_path)
+            self.training_embeddings = np.load(emb_path, allow_pickle=False)
 
     def get_centroids_data(self) -> dict:
         """Return data dict suitable for np.savez."""

{prometheusnn-0.2.1 → prometheusnn-0.2.2}/src/prometheus/ui_server.py

@@ -236,7 +236,7 @@ async def _handle_message(ws, msg: dict[str, Any]) -> None:
 try:
     from fastapi import FastAPI, WebSocket, WebSocketDisconnect
 
-    app: FastAPI | None = FastAPI(title="prometheusnn-ui", version="0.2.1")
+    app: FastAPI | None = FastAPI(title="prometheusnn-ui", version="0.2.2")
 
     @app.get("/health")
     def _health() -> dict:

prometheusnn-0.2.2/tests/test_novelty.py

@@ -0,0 +1,100 @@
+"""Tests for novelty detection."""
+
+import tempfile
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+from prometheus.inference.novelty import NoveltyDetector
+
+
+def test_fit_and_score():
+    rng = np.random.default_rng(42)
+    n = 50
+    embs = rng.standard_normal((n, 16)).astype(np.float32)
+    labels = np.array([0] * 25 + [1] * 25)
+    classes = ["classA", "classB"]
+
+    det = NoveltyDetector(z_score_threshold=2.5)
+    det.fit(embs, labels, classes)
+
+    assert len(det.centroids) == 2
+
+    # Score an in-distribution sample
+    score = det.score(embs[0], "classA", confidence=0.9)
+    assert not score.is_novel
+
+    # Score a far-out sample
+    outlier = np.ones(16, dtype=np.float32) * 100
+    score = det.score(outlier, "classA", confidence=0.3)
+    assert score.centroid_z_score > 2.5
+
+
+def test_save_load_roundtrip():
+    rng = np.random.default_rng(42)
+    embs = rng.standard_normal((30, 8)).astype(np.float32)
+    labels = np.array([0] * 15 + [1] * 15)
+    classes = ["X", "Y"]
+
+    det = NoveltyDetector()
+    det.fit(embs, labels, classes)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        det.save(Path(tmpdir))
+
+        det2 = NoveltyDetector()
+        det2.load(Path(tmpdir))
+
+        assert set(det2.centroids.keys()) == {"X", "Y"}
+        assert det2.training_embeddings is not None
+        assert det2.training_embeddings.shape == (30, 8)
+
+
+class _Payload:
+    """Unpickling this creates a file: proof that loading ran attacker code."""
+
+    def __init__(self, marker: Path):
+        self.marker = str(marker)
+
+    def __reduce__(self):
+        return (open, (self.marker, "w"))
+
+
+def test_load_never_unpickles_centroids(tmp_path):
+    marker = tmp_path / "executed"
+    np.savez(
+        tmp_path / "novelty_centroids.npz",
+        centroid_0=np.zeros(4),
+        class_names=np.array([_Payload(marker)], dtype=object),
+        mean_dists=[1.0],
+        std_dists=[1.0],
+        p95_dists=[1.0],
+        counts=[1],
+    )
+    with pytest.raises(ValueError, match="pickled objects"):
+        NoveltyDetector().load(tmp_path)
+    assert not marker.exists()
+
+
+def test_load_never_unpickles_embeddings(tmp_path):
+    marker = tmp_path / "executed"
+    det = NoveltyDetector()
+    det.fit(np.eye(4, dtype=np.float32), np.array([0, 0, 1, 1]), ["a", "b"])
+    det.save(tmp_path)
+    np.save(tmp_path / "novelty_embeddings.npy", np.array([_Payload(marker)], dtype=object))
+    with pytest.raises(ValueError):
+        NoveltyDetector().load(tmp_path)
+    assert not marker.exists()
+
+
+def test_saved_files_need_no_pickle(tmp_path):
+    rng = np.random.default_rng(0)
+    det = NoveltyDetector()
+    det.fit(rng.standard_normal((20, 4)), np.array([0] * 10 + [1] * 10), ["Café", "naïve/class"])
+    det.save(tmp_path)
+    with np.load(tmp_path / "novelty_centroids.npz", allow_pickle=False) as data:
+        assert all(data[k].dtype != object for k in data.files)
+    det2 = NoveltyDetector()
+    det2.load(tmp_path)
+    assert set(det2.centroids) == {"Café", "naïve/class"}

{prometheusnn-0.2.1 → prometheusnn-0.2.2}/uv.lock

@@ -907,7 +907,7 @@ wheels = [
 
 [[package]]
 name = "prometheusnn"
-version = "0.2.1"
+version = "0.2.2"
 source = { editable = "." }
 dependencies = [
     { name = "mcp", extra = ["cli"] },

prometheusnn-0.2.1/tests/test_novelty.py

@@ -1,50 +0,0 @@
-"""Tests for novelty detection."""
-
-import tempfile
-from pathlib import Path
-
-import numpy as np
-
-from prometheus.inference.novelty import NoveltyDetector
-
-
-def test_fit_and_score():
-    rng = np.random.default_rng(42)
-    n = 50
-    embs = rng.standard_normal((n, 16)).astype(np.float32)
-    labels = np.array([0] * 25 + [1] * 25)
-    classes = ["classA", "classB"]
-
-    det = NoveltyDetector(z_score_threshold=2.5)
-    det.fit(embs, labels, classes)
-
-    assert len(det.centroids) == 2
-
-    # Score an in-distribution sample
-    score = det.score(embs[0], "classA", confidence=0.9)
-    assert not score.is_novel
-
-    # Score a far-out sample
-    outlier = np.ones(16, dtype=np.float32) * 100
-    score = det.score(outlier, "classA", confidence=0.3)
-    assert score.centroid_z_score > 2.5
-
-
-def test_save_load_roundtrip():
-    rng = np.random.default_rng(42)
-    embs = rng.standard_normal((30, 8)).astype(np.float32)
-    labels = np.array([0] * 15 + [1] * 15)
-    classes = ["X", "Y"]
-
-    det = NoveltyDetector()
-    det.fit(embs, labels, classes)
-
-    with tempfile.TemporaryDirectory() as tmpdir:
-        det.save(Path(tmpdir))
-
-        det2 = NoveltyDetector()
-        det2.load(Path(tmpdir))
-
-        assert set(det2.centroids.keys()) == {"X", "Y"}
-        assert det2.training_embeddings is not None
-        assert det2.training_embeddings.shape == (30, 8)