predikit 0.4.6__tar.gz → 0.5.0__tar.gz

{predikit-0.4.6 → predikit-0.5.0}/CHANGELOG.md

@@ -6,7 +6,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 ## [Unreleased]
 
-## [0.4.6] - 2026-07-01
+## [0.5.0] - 2026-07-01
 
 ### Fixed
 - Improved `ModelTool` input coercion for optional scalar annotations like `float | None`, `bool | None`, `int | None`, and `str | None`

{predikit-0.4.6 → predikit-0.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: predikit
-Version: 0.4.6
+Version: 0.5.0
 Summary: Wrap any trained sklearn/XGBoost model as an LLM-callable tool: auto-generated JSON schemas, Pydantic I/O, OpenAI & LangChain ready, zero boilerplate.
 Project-URL: Homepage, https://github.com/Tejas-TA/predikit
 Project-URL: Repository, https://github.com/Tejas-TA/predikit

{predikit-0.4.6 → predikit-0.5.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "predikit"
-version = "0.4.6"
+version = "0.5.0"
 description = "Wrap any trained sklearn/XGBoost model as an LLM-callable tool: auto-generated JSON schemas, Pydantic I/O, OpenAI & LangChain ready, zero boilerplate."
 readme = "README.md"
 license = {text = "MIT"}

{predikit-0.4.6 → predikit-0.5.0}/src/predikit/__init__.py

@@ -4,4 +4,4 @@ from .registry import ToolRegistry
 from .tool import ModelTool
 
 __all__ = ["ModelTool", "ToolRegistry", "ModelEnsemble", "LowConfidenceError"]
-__version__ = "0.4.6"
+__version__ = "0.5.0"

{predikit-0.4.6 → predikit-0.5.0}/src/predikit/loaders/snowflake.py

@@ -13,26 +13,47 @@ except ImportError:
     Registry = None  # type: ignore[assignment,misc]
 
 
-class _SnowflakeShim:
-    """Wraps a Snowflake Model Registry model to look like an sklearn estimator."""
+class _SnowflakeShim:
+    """Wraps a Snowflake Model Registry model to look like an sklearn estimator."""
 
-    def __init__(
-        self, sf_model: Any, output_method: str = "predict", classes: list | None = None
-    ) -> None:
-        if not callable(getattr(sf_model, output_method, None)):
-            raise AttributeError(
-                f"Snowflake model object has no callable method {output_method!r}."
-            )
-        self._model = sf_model
-        self._method = output_method
-        if classes is not None:
-            self.classes_ = classes
-
-    def predict(self, X: Any) -> np.ndarray:
-        result = getattr(self._model, self._method)(X)
-        if hasattr(result, "to_numpy"):
-            return np.asarray(result.to_numpy()).flatten()
-        return np.asarray(result).flatten()
+    def __init__(
+        self, sf_model: Any, output_method: str = "predict", classes: list | None = None
+    ) -> None:
+        if not callable(getattr(sf_model, output_method, None)):
+            raise AttributeError(
+                f"Snowflake model object has no callable method {output_method!r}."
+            )
+        self._model = sf_model
+        self._method = output_method
+        underlying = getattr(sf_model, "_model_impl", None)
+        if underlying is not None:
+            underlying = getattr(underlying, "sklearn_model", None)
+
+        if classes is None and underlying is not None and hasattr(underlying, "classes_"):
+            classes = list(underlying.classes_)
+
+        if classes is not None:
+            self.classes_ = classes
+
+        if underlying is not None:
+            if hasattr(underlying, "feature_names_in_"):
+                self.feature_names_in_ = underlying.feature_names_in_
+            if hasattr(underlying, "n_features_in_"):
+                self.n_features_in_ = underlying.n_features_in_
+
+    def predict(self, X: Any) -> np.ndarray:
+        result = getattr(self._model, self._method)(X)
+        if hasattr(result, "to_numpy"):
+            return np.asarray(result.to_numpy()).flatten()
+        return np.asarray(result).flatten()
+
+    def predict_proba(self, X: Any) -> np.ndarray:
+        if callable(getattr(self._model, "predict_proba", None)):
+            return np.asarray(self._model.predict_proba(X))
+        raise NotImplementedError(
+            "predict_proba is not available for this Snowflake model object. "
+            "The underlying model may not expose it."
+        )
 
 
 def from_snowflake(

{predikit-0.4.6 → predikit-0.5.0}/tests/test_loaders_snowflake.py

@@ -24,7 +24,7 @@ class MemberInput(BaseModel):
 # ---------------------------------------------------------------------------
 
 
-class TestSnowflakeShim:
+class TestSnowflakeShim:
     def test_predict_calls_correct_method(self):
         mock_model = MagicMock()
         mock_model.predict.return_value = np.array([1])
@@ -41,12 +41,12 @@ class TestSnowflakeShim:
         mock_model.predict_proba.assert_called_once()
         assert result[0] == pytest.approx(0.27, abs=1e-6)
 
-    def test_classes_optional(self):
-        shim_no_classes = _SnowflakeShim(MagicMock())
-        assert not hasattr(shim_no_classes, "classes_")
-
-        shim_with_classes = _SnowflakeShim(MagicMock(), classes=[0, 1])
-        assert shim_with_classes.classes_ == [0, 1]
+    def test_classes_optional(self):
+        shim_no_classes = _SnowflakeShim(MagicMock(spec=["predict"]))
+        assert not hasattr(shim_no_classes, "classes_")
+
+        shim_with_classes = _SnowflakeShim(MagicMock(spec=["predict"]), classes=[0, 1])
+        assert shim_with_classes.classes_ == [0, 1]
 
     def test_predict_flattens_dataframe_result(self):
         import pandas as pd
@@ -58,13 +58,24 @@ class TestSnowflakeShim:
         assert result.ndim == 1
         assert result[0] == 1
 
-    def test_predict_flattens_numpy_result(self):
-        mock_model = MagicMock()
-        mock_model.predict.return_value = np.array([[0.73]])
-        shim = _SnowflakeShim(mock_model)
-        result = shim.predict(np.array([[1.0]]))
-        assert result.ndim == 1
-        assert result[0] == pytest.approx(0.73)
+    def test_predict_flattens_numpy_result(self):
+        mock_model = MagicMock()
+        mock_model.predict.return_value = np.array([[0.73]])
+        shim = _SnowflakeShim(mock_model)
+        result = shim.predict(np.array([[1.0]]))
+        assert result.ndim == 1
+        assert result[0] == pytest.approx(0.73)
+
+    def test_predict_proba_delegates_to_underlying_model(self):
+        mock_model = MagicMock()
+        mock_model.predict_proba.return_value = np.array([[0.2, 0.8]])
+        shim = _SnowflakeShim(mock_model)
+
+        result = shim.predict_proba(np.array([[1.0, 2.0, 3.0]]))
+
+        mock_model.predict_proba.assert_called_once()
+        assert result.shape == (1, 2)
+        assert result[0, 1] == pytest.approx(0.8)
 
 
 # ---------------------------------------------------------------------------
@@ -80,7 +91,7 @@ def _make_mock_session_and_model(return_value):
     return mock_session, mock_sf_model
 
 
-class TestFromSnowflake:
+class TestFromSnowflake:
     def test_roundtrip_invoke(self):
         mock_session, mock_sf_model = _make_mock_session_and_model(np.array([1]))
 
@@ -123,10 +134,10 @@ class TestFromSnowflake:
             mock_registry.get_model.assert_called_once_with("VACATION_CHURN")
             mock_registry.get_model.return_value.version.assert_called_once_with("V3")
 
-    def test_custom_output_method_forwarded(self):
-        mock_session = MagicMock()
-        mock_sf_model = MagicMock()
-        mock_sf_model.score.return_value = np.array([0.87])
+    def test_custom_output_method_forwarded(self):
+        mock_session = MagicMock()
+        mock_sf_model = MagicMock()
+        mock_sf_model.score.return_value = np.array([0.87])
 
         with patch("predikit.loaders.snowflake.Registry") as MockRegistry:
             MockRegistry.return_value.get_model.return_value.version.return_value = mock_sf_model
@@ -143,9 +154,38 @@ class TestFromSnowflake:
                 output_method="score",
             )
 
-        result = tool.invoke({"tenure_months": 12.0, "trips_last_year": 5.0, "avg_spend": 300.0})
-        assert result["score"] == pytest.approx(0.87, abs=1e-6)
-        mock_sf_model.score.assert_called_once()
+        result = tool.invoke({"tenure_months": 12.0, "trips_last_year": 5.0, "avg_spend": 300.0})
+        assert result["score"] == pytest.approx(0.87, abs=1e-6)
+        mock_sf_model.score.assert_called_once()
+
+    def test_confidence_routing_uses_predict_proba(self):
+        mock_session = MagicMock()
+        mock_sf_model = MagicMock()
+        mock_sf_model.predict.return_value = np.array([1])
+        mock_sf_model.predict_proba.return_value = np.array([[0.6, 0.4]])
+        mock_sf_model.classes_ = [0, 1]
+
+        with patch("predikit.loaders.snowflake.Registry") as MockRegistry:
+            MockRegistry.return_value.get_model.return_value.version.return_value = mock_sf_model
+
+            tool = from_snowflake(
+                session=mock_session,
+                model_name="CHURN",
+                model_version="V1",
+                name="churn",
+                description="Churn.",
+                input_schema=MemberInput,
+                output_name="churn_class",
+                output_description="Churn class",
+                confidence_threshold=0.9,
+                on_low_confidence="warn",
+            )
+
+        result = tool.invoke({"tenure_months": 12.0, "trips_last_year": 5.0, "avg_spend": 300.0})
+        assert result["churn_class"] == 1
+        assert result["_low_confidence"] is True
+        assert result["_confidence"] == pytest.approx(0.6)
+        mock_sf_model.predict_proba.assert_called_once()
 
     def test_model_tool_kwargs_forwarded(self):
         mock_session, mock_sf_model = _make_mock_session_and_model(np.array([1]))

{predikit-0.4.6 → predikit-0.5.0}/uv.lock

@@ -560,7 +560,7 @@ wheels = [
 
 [[package]]
 name = "colorama"
-version = "0.4.6"
+version = "0.5.0"
 source = { registry = "https://pypi.org/simple" }
 sdist = { url = "https://files.pythonhosted.org/packages/d8/53/6f443c9a4a8358a93a6792e2acffb9d9d5cb0a5cfd8802644b7b1c9a02e4/colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44", size = 27697, upload-time = "2022-10-25T02:36:22.414Z" }
 wheels = [