predikit 0.6.0__tar.gz → 0.6.2__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (62) hide show
  1. {predikit-0.6.0 → predikit-0.6.2}/.github/workflows/test.yml +13 -13
  2. predikit-0.6.2/.gitignore +33 -0
  3. {predikit-0.6.0 → predikit-0.6.2}/CHANGELOG.md +49 -19
  4. {predikit-0.6.0 → predikit-0.6.2}/CLAUDE.md +4 -4
  5. {predikit-0.6.0 → predikit-0.6.2}/CONTRIBUTING.md +2 -2
  6. {predikit-0.6.0 → predikit-0.6.2}/PKG-INFO +4 -2
  7. {predikit-0.6.0 → predikit-0.6.2}/pyproject.toml +9 -7
  8. predikit-0.6.2/src/predikit/__init__.py +16 -0
  9. {predikit-0.6.0 → predikit-0.6.2}/src/predikit/cli.py +30 -11
  10. {predikit-0.6.0 → predikit-0.6.2}/src/predikit/coerce.py +4 -0
  11. {predikit-0.6.0 → predikit-0.6.2}/src/predikit/loaders/snowflake.py +46 -49
  12. predikit-0.6.2/src/predikit/mcp.py +115 -0
  13. {predikit-0.6.0 → predikit-0.6.2}/src/predikit/registry.py +19 -7
  14. predikit-0.6.2/tests/__init__.py +0 -0
  15. predikit-0.6.2/tests/test_cli.py +219 -0
  16. {predikit-0.6.0 → predikit-0.6.2}/tests/test_coerce.py +11 -0
  17. predikit-0.6.2/tests/test_loaders_snowflake.py +250 -0
  18. {predikit-0.6.0 → predikit-0.6.2}/tests/test_logging.py +52 -39
  19. predikit-0.6.2/tests/test_mcp.py +158 -0
  20. {predikit-0.6.0 → predikit-0.6.2}/tests/test_registry.py +15 -0
  21. predikit-0.6.0/.claude/settings.local.json +0 -36
  22. predikit-0.6.0/.gitignore +0 -14
  23. predikit-0.6.0/src/predikit/__init__.py +0 -8
  24. predikit-0.6.0/src/predikit/mcp.py +0 -68
  25. predikit-0.6.0/tests/test_cli.py +0 -104
  26. predikit-0.6.0/tests/test_loaders_snowflake.py +0 -240
  27. predikit-0.6.0/tests/test_mcp.py +0 -82
  28. predikit-0.6.0/uv.lock +0 -5586
  29. {predikit-0.6.0 → predikit-0.6.2}/.gitattributes +0 -0
  30. {predikit-0.6.0 → predikit-0.6.2}/.github/ISSUE_TEMPLATE/bug_report.md +0 -0
  31. {predikit-0.6.0 → predikit-0.6.2}/.github/ISSUE_TEMPLATE/feature_request.md +0 -0
  32. {predikit-0.6.0 → predikit-0.6.2}/.github/workflows/publish.yml +0 -0
  33. {predikit-0.6.0 → predikit-0.6.2}/.pre-commit-config.yaml +0 -0
  34. {predikit-0.6.0 → predikit-0.6.2}/LICENSE +0 -0
  35. {predikit-0.6.0 → predikit-0.6.2}/README.md +0 -0
  36. {predikit-0.6.0 → predikit-0.6.2}/docs/logo.gif +0 -0
  37. {predikit-0.6.0 → predikit-0.6.2}/docs/logo.png +0 -0
  38. {predikit-0.6.0 → predikit-0.6.2}/examples/01_basic_sklearn.py +0 -0
  39. {predikit-0.6.0 → predikit-0.6.2}/examples/02_xgboost_regression.py +0 -0
  40. {predikit-0.6.0 → predikit-0.6.2}/examples/03_orlando_real_estate.py +0 -0
  41. {predikit-0.6.0 → predikit-0.6.2}/examples/04_confidence_routing.py +0 -0
  42. {predikit-0.6.0 → predikit-0.6.2}/examples/05_multi_model_ensemble.py +0 -0
  43. {predikit-0.6.0 → predikit-0.6.2}/examples/06_mlflow_loader.py +0 -0
  44. {predikit-0.6.0 → predikit-0.6.2}/examples/07_snowflake_loader.py +0 -0
  45. {predikit-0.6.0 → predikit-0.6.2}/src/predikit/ensemble.py +0 -0
  46. {predikit-0.6.0 → predikit-0.6.2}/src/predikit/exceptions.py +0 -0
  47. {predikit-0.6.0 → predikit-0.6.2}/src/predikit/exporters/__init__.py +0 -0
  48. {predikit-0.6.0 → predikit-0.6.2}/src/predikit/exporters/langchain.py +0 -0
  49. {predikit-0.6.0 → predikit-0.6.2}/src/predikit/exporters/openai.py +0 -0
  50. {predikit-0.6.0 → predikit-0.6.2}/src/predikit/introspect.py +0 -0
  51. {predikit-0.6.0 → predikit-0.6.2}/src/predikit/loaders/__init__.py +0 -0
  52. {predikit-0.6.0 → predikit-0.6.2}/src/predikit/loaders/mlflow.py +0 -0
  53. /predikit-0.6.0/tests/__init__.py → /predikit-0.6.2/src/predikit/py.typed +0 -0
  54. {predikit-0.6.0 → predikit-0.6.2}/src/predikit/tool.py +0 -0
  55. {predikit-0.6.0 → predikit-0.6.2}/tests/test_confidence.py +0 -0
  56. {predikit-0.6.0 → predikit-0.6.2}/tests/test_ensemble.py +0 -0
  57. {predikit-0.6.0 → predikit-0.6.2}/tests/test_exporters_langchain.py +0 -0
  58. {predikit-0.6.0 → predikit-0.6.2}/tests/test_exporters_openai.py +0 -0
  59. {predikit-0.6.0 → predikit-0.6.2}/tests/test_introspect.py +0 -0
  60. {predikit-0.6.0 → predikit-0.6.2}/tests/test_loaders_mlflow.py +0 -0
  61. {predikit-0.6.0 → predikit-0.6.2}/tests/test_tool.py +0 -0
  62. {predikit-0.6.0 → predikit-0.6.2}/tests/test_weighted_ensemble.py +0 -0
@@ -23,24 +23,24 @@ jobs:
23
23
  run: pip install -e ".[dev]"
24
24
 
25
25
  - name: Ruff lint
26
- run: ruff check src/ tests/
26
+ run: ruff check src/ tests/ examples/
27
27
 
28
28
  - name: Ruff format check
29
- run: ruff format --check src/ tests/
30
-
31
- - name: Mypy
32
- run: mypy src/predikit
33
-
34
- - name: Build package
35
- run: |
36
- python -m pip install hatchling
37
- python -m hatchling build
38
-
39
- test:
29
+ run: ruff format --check src/ tests/ examples/
30
+
31
+ - name: Mypy
32
+ run: mypy src/predikit
33
+
34
+ - name: Build package
35
+ run: |
36
+ python -m pip install hatchling
37
+ python -m hatchling build
38
+
39
+ test:
40
40
  runs-on: ubuntu-latest
41
41
  strategy:
42
42
  matrix:
43
- python-version: ["3.10", "3.11", "3.12"]
43
+ python-version: ["3.10", "3.11", "3.12", "3.13"]
44
44
 
45
45
  steps:
46
46
  - uses: actions/checkout@v4
@@ -0,0 +1,33 @@
1
+ # Byte-compiled / cache
2
+ __pycache__/
3
+ *.py[cod]
4
+
5
+ # Distribution / packaging
6
+ build/
7
+ dist/
8
+ *.egg
9
+ *.egg-info/
10
+
11
+ # Virtual environments
12
+ .venv/
13
+ venv/
14
+
15
+ # Tooling caches
16
+ .mypy_cache/
17
+ .pytest_cache/
18
+ .ruff_cache/
19
+
20
+ # Test coverage
21
+ .coverage
22
+ coverage.xml
23
+ htmlcov/
24
+
25
+ # Serialized models written by examples and CLI tests
26
+ *.joblib
27
+ *.pkl
28
+
29
+ # MLflow local tracking store
30
+ mlruns/
31
+
32
+ # Per-developer editor and agent settings
33
+ .claude/settings.local.json
@@ -4,25 +4,55 @@ All notable changes to this project will be documented in this file.
4
4
 
5
5
  The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
6
6
 
7
- ## [Unreleased]
8
-
9
- ## [0.6.0] - 2026-08-03
10
-
11
- ### Added
12
- - MCP server integration via `create_mcp_server()` and `predikit serve`.
13
- - Optional `mcp` dependency for exposing `ToolRegistry` tools to MCP-compatible clients.
14
-
15
- ## [0.5.2] - 2026-08-01
16
-
17
- ### Fixed
18
- - Repaired the development lockfile so `uv` can resolve the pinned Colorama artifact.
19
- - Added repository-wide LF line-ending normalization so Ruff formatting is stable across platforms.
20
-
21
- ### Added
22
- - Regression coverage for invalid Snowflake output methods and duplicate registry names.
23
- - CI package-build verification to catch release artifacts that cannot be built.
24
-
25
- ## [0.5.1] - 2026-07-21
7
+ ## [Unreleased]
8
+
9
+ ## [0.6.2] - 2026-08-22
10
+
11
+ ### Added
12
+ - `py.typed` marker, so type checkers actually consume the annotations. Every module is annotated and mypy-clean, and the package declares the `Typing :: Typed` classifier, but without the PEP 561 marker downstream consumers got no type information at all.
13
+
14
+ ### Changed
15
+ - `from_snowflake()` imports `snowflake.ml.registry` inside the function body instead of at module level, matching `from_mlflow()` and the project's own rule for optional dependencies. The module-level `Registry = None` sentinel is gone.
16
+ - CI lints and format-checks `examples/` alongside `src/` and `tests/`.
17
+
18
+ ### Removed
19
+ - `uv.lock`. Nothing consumed it — CI installs with `pip install -e ".[dev]"` — and a library should not pin its consumers' transitive dependencies.
20
+ - `.claude/settings.local.json` from version control. It holds per-developer machine paths.
21
+
22
+ ## [0.6.1] - 2026-08-15
23
+
24
+ ### Fixed
25
+ - MCP tools now advertise an `outputSchema` and return `structuredContent`. The generated callable carried no return annotation, so FastMCP emitted `outputSchema: null` and every result came back as unstructured text. Output fields are named from the tool's `output_name` / `output_description`, and a `collect` ensemble declares each member's output.
26
+ - Low-confidence results keep their `_confidence` and `_low_confidence` keys when served over MCP. The output model allows extra keys; a strict one validates them away silently.
27
+ - `predikit serve` no longer disguises unrelated failures as a bad registry target. Module import and attribute lookup are handled separately, and errors raised while building or running the server propagate with their own traceback instead of being reported as `Could not load registry target`.
28
+ - `coerce_value()` converts `True` / `False` to `1` / `0` for `int` fields. `isinstance(True, int)` is true in Python, so bools previously passed through unconverted.
29
+
30
+ ### Added
31
+ - `--host` and `--port` options on `predikit serve`, plus matching `host` / `port` arguments on `create_mcp_server()`. The `streamable-http` transport was previously stuck on the SDK default of `127.0.0.1:8000`.
32
+ - `ToolRegistry.items()`, `.names()`, `__iter__`, and `__len__` as the public way to enumerate a registry. `create_mcp_server()` used the private `_tools` / `_ensembles` attributes; `to_openai()` and `to_langchain()` now share the same accessor.
33
+ - Python 3.13 in the CI matrix and the package classifiers.
34
+
35
+ ### Changed
36
+ - `mcp` is now a dev dependency, so the MCP integration is tested against the real SDK. The previous tests substituted a fake `FastMCP` and reported full coverage of `mcp.py` while never exercising the SDK — which is how the missing output schema shipped.
37
+ - Test coverage for `predikit serve`, which previously had none.
38
+
39
+ ## [0.6.0] - 2026-08-03
40
+
41
+ ### Added
42
+ - MCP server integration via `create_mcp_server()` and `predikit serve`.
43
+ - Optional `mcp` dependency for exposing `ToolRegistry` tools to MCP-compatible clients.
44
+
45
+ ## [0.5.2] - 2026-08-01
46
+
47
+ ### Fixed
48
+ - Repaired the development lockfile so `uv` can resolve the pinned Colorama artifact.
49
+ - Added repository-wide LF line-ending normalization so Ruff formatting is stable across platforms.
50
+
51
+ ### Added
52
+ - Regression coverage for invalid Snowflake output methods and duplicate registry names.
53
+ - CI package-build verification to catch release artifacts that cannot be built.
54
+
55
+ ## [0.5.1] - 2026-07-21
26
56
 
27
57
  ### Added
28
58
  - Test coverage for the LangChain exporter (`test_exporters_langchain.py`) — previously untested
@@ -35,14 +35,14 @@ pip install -e ".[dev]"
35
35
  pytest --cov=src/predikit --cov-report=term-missing
36
36
 
37
37
  # Lint (check only)
38
- ruff check src/ tests/
38
+ ruff check src/ tests/ examples/
39
39
 
40
40
  # Format (check only)
41
- ruff format --check src/ tests/
41
+ ruff format --check src/ tests/ examples/
42
42
 
43
43
  # Lint + format (auto-fix)
44
- ruff check --fix src/ tests/
45
- ruff format src/ tests/
44
+ ruff check --fix src/ tests/ examples/
45
+ ruff format src/ tests/ examples/
46
46
 
47
47
  # Type check
48
48
  mypy src/predikit
@@ -30,8 +30,8 @@ This project uses [ruff](https://docs.astral.sh/ruff/) for linting and formattin
30
30
 
31
31
  ```bash
32
32
  # Check and auto-fix
33
- ruff check --fix src/ tests/
34
- ruff format src/ tests/
33
+ ruff check --fix src/ tests/ examples/
34
+ ruff format src/ tests/ examples/
35
35
 
36
36
  # Type check
37
37
  mypy src/predikit
@@ -1,6 +1,6 @@
1
- Metadata-Version: 2.4
1
+ Metadata-Version: 2.5
2
2
  Name: predikit
3
- Version: 0.6.0
3
+ Version: 0.6.2
4
4
  Summary: Wrap any trained sklearn/XGBoost model as an LLM-callable tool: auto-generated JSON schemas, Pydantic I/O, OpenAI & LangChain ready, zero boilerplate.
5
5
  Project-URL: Homepage, https://github.com/Tejas-TA/predikit
6
6
  Project-URL: Repository, https://github.com/Tejas-TA/predikit
@@ -19,6 +19,7 @@ Classifier: Programming Language :: Python :: 3
19
19
  Classifier: Programming Language :: Python :: 3.10
20
20
  Classifier: Programming Language :: Python :: 3.11
21
21
  Classifier: Programming Language :: Python :: 3.12
22
+ Classifier: Programming Language :: Python :: 3.13
22
23
  Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
23
24
  Classifier: Topic :: Software Development :: Libraries :: Python Modules
24
25
  Classifier: Typing :: Typed
@@ -33,6 +34,7 @@ Provides-Extra: dev
33
34
  Requires-Dist: click>=8.0; extra == 'dev'
34
35
  Requires-Dist: joblib>=1.2; extra == 'dev'
35
36
  Requires-Dist: langchain-core>=0.1; extra == 'dev'
37
+ Requires-Dist: mcp<2,>=1.28; extra == 'dev'
36
38
  Requires-Dist: mlflow>=2.0; extra == 'dev'
37
39
  Requires-Dist: mypy>=1.10; extra == 'dev'
38
40
  Requires-Dist: pandas>=1.5; extra == 'dev'
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "predikit"
7
- version = "0.6.0"
7
+ version = "0.6.2"
8
8
  description = "Wrap any trained sklearn/XGBoost model as an LLM-callable tool: auto-generated JSON schemas, Pydantic I/O, OpenAI & LangChain ready, zero boilerplate."
9
9
  readme = "README.md"
10
10
  license = {text = "MIT"}
@@ -25,6 +25,7 @@ classifiers = [
25
25
  "Programming Language :: Python :: 3.10",
26
26
  "Programming Language :: Python :: 3.11",
27
27
  "Programming Language :: Python :: 3.12",
28
+ "Programming Language :: Python :: 3.13",
28
29
  "Topic :: Scientific/Engineering :: Artificial Intelligence",
29
30
  "Intended Audience :: Developers",
30
31
  "Intended Audience :: Science/Research",
@@ -38,9 +39,9 @@ dependencies = [
38
39
  ]
39
40
 
40
41
  [project.optional-dependencies]
41
- langchain = ["langchain-core>=0.1"]
42
- mcp = ["mcp>=1.28,<2"]
43
- xgboost = ["xgboost>=1.7"]
42
+ langchain = ["langchain-core>=0.1"]
43
+ mcp = ["mcp>=1.28,<2"]
44
+ xgboost = ["xgboost>=1.7"]
44
45
  mlflow = ["mlflow>=2.0"]
45
46
  snowflake = ["snowflake-ml-python>=1.0"]
46
47
  cli = ["click>=8.0", "joblib>=1.2"]
@@ -50,6 +51,7 @@ dev = [
50
51
  "xgboost>=1.7",
51
52
  "langchain-core>=0.1",
52
53
  "pandas>=1.5",
54
+ "mcp>=1.28,<2",
53
55
  "mlflow>=2.0",
54
56
  "click>=8.0",
55
57
  "joblib>=1.2",
@@ -78,9 +80,9 @@ markers = [
78
80
  "integration: tests that exercise slower external runtimes or registries",
79
81
  ]
80
82
 
81
- [tool.ruff]
82
- line-length = 100
83
- target-version = "py310"
83
+ [tool.ruff]
84
+ line-length = 100
85
+ target-version = "py310"
84
86
 
85
87
  [tool.ruff.lint]
86
88
  select = ["E", "F", "I", "UP", "B", "SIM"]
@@ -0,0 +1,16 @@
1
+ """Wrap fitted scikit-learn / XGBoost models as LLM-callable tools."""
2
+
3
+ from .ensemble import ModelEnsemble
4
+ from .exceptions import LowConfidenceError
5
+ from .mcp import create_mcp_server
6
+ from .registry import ToolRegistry
7
+ from .tool import ModelTool
8
+
9
+ __all__ = [
10
+ "LowConfidenceError",
11
+ "ModelEnsemble",
12
+ "ModelTool",
13
+ "ToolRegistry",
14
+ "create_mcp_server",
15
+ ]
16
+ __version__ = "0.6.2"
@@ -60,7 +60,7 @@ if _CLICK_AVAILABLE:
60
60
  click.echo(f"Classes: {meta['classes']}")
61
61
 
62
62
  if meta["feature_names"]:
63
- fields: dict[str, Any] = {f: (float, ...) for f in meta["feature_names"]}
63
+ fields: dict[str, Any] = dict.fromkeys(meta["feature_names"], (float, ...))
64
64
  input_schema = create_model("Input", **fields)
65
65
  tool = ModelTool(
66
66
  model=model,
@@ -85,7 +85,19 @@ if _CLICK_AVAILABLE:
85
85
  show_default=True,
86
86
  help="MCP transport to run.",
87
87
  )
88
- def serve(registry_target: str, name: str, transport: str) -> None:
88
+ @click.option(
89
+ "--host", default=None, help="Bind host for HTTP transports. [default: 127.0.0.1]"
90
+ )
91
+ @click.option(
92
+ "--port", type=int, default=None, help="Bind port for HTTP transports. [default: 8000]"
93
+ )
94
+ def serve(
95
+ registry_target: str,
96
+ name: str,
97
+ transport: str,
98
+ host: str | None,
99
+ port: int | None,
100
+ ) -> None:
89
101
  """Serve a ToolRegistry from MODULE:ATTRIBUTE over MCP.
90
102
 
91
103
  ATTRIBUTE may be a ToolRegistry instance or a zero-argument factory.
@@ -93,22 +105,29 @@ if _CLICK_AVAILABLE:
93
105
  if ":" not in registry_target:
94
106
  raise click.ClickException("REGISTRY_TARGET must use the form MODULE:ATTRIBUTE")
95
107
  module_name, attribute_name = registry_target.split(":", 1)
108
+
109
+ # Keep this block narrow: it resolves the target only. Errors raised while
110
+ # building or running the server must surface with their own traceback
111
+ # rather than be reported as a bad registry target.
96
112
  try:
97
113
  module = importlib.import_module(module_name)
98
- registry = getattr(module, attribute_name)
99
- if callable(registry):
100
- registry = registry()
101
- from .mcp import create_mcp_server
102
-
103
- server = create_mcp_server(registry, name=name)
104
- server.run(transport=transport)
105
114
  except ImportError as err:
106
- raise click.ClickException(str(err)) from err
115
+ raise click.ClickException(f"Could not import module '{module_name}': {err}") from err
116
+ try:
117
+ registry = getattr(module, attribute_name)
107
118
  except AttributeError as err:
108
119
  raise click.ClickException(
109
- f"Could not load registry target '{registry_target}'"
120
+ f"Module '{module_name}' has no attribute '{attribute_name}'"
110
121
  ) from err
111
122
 
123
+ if callable(registry):
124
+ registry = registry()
125
+
126
+ from .mcp import create_mcp_server
127
+
128
+ server = create_mcp_server(registry, name=name, host=host, port=port)
129
+ server.run(transport=transport)
130
+
112
131
  else:
113
132
 
114
133
  def cli() -> None: # type: ignore[misc]
@@ -8,6 +8,10 @@ _BOOL_FALSE = {"false", "0", "no", "off"}
8
8
 
9
9
  def coerce_value(value: Any, target_type: type) -> Any:
10
10
  """Coerce a single value to target_type with LLM-friendly string handling."""
11
+ if target_type is int and isinstance(value, bool):
12
+ # isinstance(True, int) is True, so bools would otherwise pass through unconverted.
13
+ return int(value)
14
+
11
15
  if isinstance(value, target_type):
12
16
  return value
13
17
 
@@ -7,53 +7,48 @@ from pydantic import BaseModel
7
7
 
8
8
  from ..tool import ModelTool
9
9
 
10
- try:
11
- from snowflake.ml.registry import Registry
12
- except ImportError:
13
- Registry = None # type: ignore[assignment,misc]
14
-
15
-
16
- class _SnowflakeShim:
17
- """Wrap a Snowflake Model Registry model as an sklearn estimator."""
18
-
19
- def __init__(
20
- self, sf_model: Any, output_method: str = "predict", classes: list | None = None
21
- ) -> None:
22
- if not callable(getattr(sf_model, output_method, None)):
23
- raise AttributeError(
24
- f"Snowflake model object has no callable method {output_method!r}."
25
- )
26
- self._model = sf_model
27
- self._method = output_method
28
- underlying = getattr(sf_model, "_model_impl", None)
29
- if underlying is not None:
30
- underlying = getattr(underlying, "sklearn_model", None)
31
-
32
- if classes is None and underlying is not None and hasattr(underlying, "classes_"):
33
- classes = list(underlying.classes_)
34
-
35
- if classes is not None:
36
- self.classes_ = classes
37
-
38
- if underlying is not None:
39
- if hasattr(underlying, "feature_names_in_"):
40
- self.feature_names_in_ = underlying.feature_names_in_
41
- if hasattr(underlying, "n_features_in_"):
42
- self.n_features_in_ = underlying.n_features_in_
43
-
44
- def predict(self, X: Any) -> np.ndarray:
45
- result = getattr(self._model, self._method)(X)
46
- if hasattr(result, "to_numpy"):
47
- return np.asarray(result.to_numpy()).flatten()
48
- return np.asarray(result).flatten()
49
-
50
- def predict_proba(self, X: Any) -> np.ndarray:
51
- if callable(getattr(self._model, "predict_proba", None)):
52
- return np.asarray(self._model.predict_proba(X))
53
- raise NotImplementedError(
54
- "predict_proba is not available for this Snowflake model object. "
55
- "The underlying model may not expose it."
56
- )
10
+
11
+ class _SnowflakeShim:
12
+ """Wrap a Snowflake Model Registry model as an sklearn estimator."""
13
+
14
+ def __init__(
15
+ self, sf_model: Any, output_method: str = "predict", classes: list | None = None
16
+ ) -> None:
17
+ if not callable(getattr(sf_model, output_method, None)):
18
+ raise AttributeError(
19
+ f"Snowflake model object has no callable method {output_method!r}."
20
+ )
21
+ self._model = sf_model
22
+ self._method = output_method
23
+ underlying = getattr(sf_model, "_model_impl", None)
24
+ if underlying is not None:
25
+ underlying = getattr(underlying, "sklearn_model", None)
26
+
27
+ if classes is None and underlying is not None and hasattr(underlying, "classes_"):
28
+ classes = list(underlying.classes_)
29
+
30
+ if classes is not None:
31
+ self.classes_ = classes
32
+
33
+ if underlying is not None:
34
+ if hasattr(underlying, "feature_names_in_"):
35
+ self.feature_names_in_ = underlying.feature_names_in_
36
+ if hasattr(underlying, "n_features_in_"):
37
+ self.n_features_in_ = underlying.n_features_in_
38
+
39
+ def predict(self, X: Any) -> np.ndarray:
40
+ result = getattr(self._model, self._method)(X)
41
+ if hasattr(result, "to_numpy"):
42
+ return np.asarray(result.to_numpy()).flatten()
43
+ return np.asarray(result).flatten()
44
+
45
+ def predict_proba(self, X: Any) -> np.ndarray:
46
+ if callable(getattr(self._model, "predict_proba", None)):
47
+ return np.asarray(self._model.predict_proba(X))
48
+ raise NotImplementedError(
49
+ "predict_proba is not available for this Snowflake model object. "
50
+ "The underlying model may not expose it."
51
+ )
57
52
 
58
53
 
59
54
  def from_snowflake(
@@ -87,11 +82,13 @@ def from_snowflake(
87
82
  Returns:
88
83
  A fully configured :class:`~predikit.ModelTool`.
89
84
  """
90
- if Registry is None:
85
+ try:
86
+ from snowflake.ml.registry import Registry
87
+ except ImportError as exc:
91
88
  raise ImportError(
92
89
  "snowflake-ml-python is required for from_snowflake(). "
93
90
  "Install it with: pip install predikit[snowflake]"
94
- )
91
+ ) from exc
95
92
 
96
93
  registry = Registry(session=session)
97
94
  sf_model = registry.get_model(model_name).version(model_version)
@@ -0,0 +1,115 @@
1
+ """MCP server integration for predikit tools."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import inspect
6
+ from typing import Any, cast
7
+
8
+ from pydantic import ConfigDict, Field, create_model
9
+
10
+ from .registry import ToolRegistry
11
+
12
+
13
+ def create_mcp_server(
14
+ registry: ToolRegistry,
15
+ name: str = "predikit",
16
+ host: str | None = None,
17
+ port: int | None = None,
18
+ ) -> Any:
19
+ """Create an MCP server exposing every tool in ``registry``.
20
+
21
+ The MCP dependency is optional. Install it with ``pip install predikit[mcp]``.
22
+ The returned server uses the SDK's standard transports, including ``stdio``.
23
+
24
+ ``host`` and ``port`` configure the HTTP transports; leave them unset to keep
25
+ the MCP SDK defaults (``127.0.0.1:8000``).
26
+ """
27
+ try:
28
+ from mcp.server.fastmcp import FastMCP
29
+ except ImportError as err:
30
+ raise ImportError(
31
+ "The MCP integration requires the optional dependency. "
32
+ "Install it with: pip install predikit[mcp]"
33
+ ) from err
34
+
35
+ settings: dict[str, Any] = {"json_response": True}
36
+ if host is not None:
37
+ settings["host"] = host
38
+ if port is not None:
39
+ settings["port"] = port
40
+
41
+ server = FastMCP(name, **settings)
42
+ for tool in registry.items():
43
+ server.add_tool(
44
+ _make_mcp_callable(tool),
45
+ name=tool.name,
46
+ description=tool.description,
47
+ )
48
+ return server
49
+
50
+
51
+ def _output_fields(tool: Any) -> list[tuple[str, str]]:
52
+ """Return the (name, description) pairs a tool's result dict is keyed by."""
53
+ tools = getattr(tool, "tools", None)
54
+ if tools is not None and getattr(tool, "strategy", None) == "collect":
55
+ # A collect ensemble merges every member's output into one dict.
56
+ return [(t.output_name, t.output_description) for t in tools]
57
+ return [(tool.output_name, tool.output_description)]
58
+
59
+
60
+ def _make_output_model(tool: Any) -> Any:
61
+ """Build the return annotation used to advertise a tool's output schema.
62
+
63
+ ``extra="allow"`` matters: it keeps the ``_confidence`` / ``_low_confidence``
64
+ keys that a low-confidence result carries. A model without it validates them
65
+ away, so MCP clients would silently lose them.
66
+ """
67
+ fields = _output_fields(tool)
68
+ if not all(name.isidentifier() for name, _ in fields):
69
+ # Output names are user-supplied; fall back to a permissive mapping.
70
+ return dict[str, Any]
71
+ definitions: dict[str, Any] = {
72
+ name: (Any, Field(..., description=desc)) for name, desc in fields
73
+ }
74
+ return create_model(
75
+ f"{tool.name}_output",
76
+ __config__=ConfigDict(extra="allow"),
77
+ **definitions,
78
+ )
79
+
80
+
81
+ def _make_mcp_callable(tool: Any) -> Any:
82
+ """Build a callable whose signature mirrors a tool's Pydantic input model."""
83
+
84
+ async def invoke(**kwargs: Any) -> dict:
85
+ return cast(dict, await tool.ainvoke(kwargs))
86
+
87
+ parameters = []
88
+ annotations: dict[str, Any] = {}
89
+ for name, field in tool.input_schema.model_fields.items():
90
+ annotation = field.annotation if field.annotation is not None else Any
91
+ annotations[name] = annotation
92
+ default = (
93
+ inspect.Parameter.empty
94
+ if field.is_required()
95
+ else field.get_default(call_default_factory=True)
96
+ )
97
+ parameters.append(
98
+ inspect.Parameter(
99
+ name,
100
+ inspect.Parameter.KEYWORD_ONLY,
101
+ annotation=annotation,
102
+ default=default,
103
+ )
104
+ )
105
+ # FastMCP derives outputSchema from the return annotation. A bare ``dict`` is
106
+ # not serializable for structured output, so it must be a concrete model.
107
+ output_model = _make_output_model(tool)
108
+ annotations["return"] = output_model
109
+ invoke.__signature__ = inspect.Signature( # type: ignore[attr-defined]
110
+ parameters, return_annotation=output_model
111
+ )
112
+ invoke.__annotations__ = annotations
113
+ invoke.__name__ = tool.name
114
+ invoke.__doc__ = tool.description
115
+ return invoke
@@ -1,5 +1,7 @@
1
1
  from __future__ import annotations
2
2
 
3
+ from collections.abc import Iterator
4
+
3
5
  from .ensemble import ModelEnsemble
4
6
  from .tool import ModelTool
5
7
 
@@ -26,15 +28,25 @@ class ToolRegistry:
26
28
  return self._tools[name]
27
29
  if name in self._ensembles:
28
30
  return self._ensembles[name]
29
- available = list(self._tools) + list(self._ensembles)
31
+ available = self.names()
30
32
  raise KeyError(f"No tool or ensemble named '{name}'. Available: {available}")
31
33
 
34
+ def items(self) -> list[RegistryItem]:
35
+ """Return every tool followed by every ensemble, in registration order."""
36
+ return list(self._tools.values()) + list(self._ensembles.values())
37
+
38
+ def names(self) -> list[str]:
39
+ """Return the name of every registered tool and ensemble."""
40
+ return [item.name for item in self.items()]
41
+
42
+ def __iter__(self) -> Iterator[RegistryItem]:
43
+ return iter(self.items())
44
+
45
+ def __len__(self) -> int:
46
+ return len(self._tools) + len(self._ensembles)
47
+
32
48
  def to_openai(self) -> list[dict]:
33
- return [t.to_openai() for t in self._tools.values()] + [
34
- e.to_openai() for e in self._ensembles.values()
35
- ]
49
+ return [item.to_openai() for item in self.items()]
36
50
 
37
51
  def to_langchain(self) -> list:
38
- return [t.to_langchain() for t in self._tools.values()] + [
39
- e.to_langchain() for e in self._ensembles.values()
40
- ]
52
+ return [item.to_langchain() for item in self.items()]
File without changes