predikit 0.5.2__tar.gz → 0.6.1__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 (56) hide show
  1. {predikit-0.5.2 → predikit-0.6.1}/.claude/settings.local.json +13 -1
  2. {predikit-0.5.2 → predikit-0.6.1}/.github/workflows/test.yml +1 -1
  3. {predikit-0.5.2 → predikit-0.6.1}/CHANGELOG.md +23 -0
  4. {predikit-0.5.2 → predikit-0.6.1}/PKG-INFO +21 -3
  5. {predikit-0.5.2 → predikit-0.6.1}/README.md +15 -1
  6. {predikit-0.5.2 → predikit-0.6.1}/pyproject.toml +4 -1
  7. {predikit-0.5.2 → predikit-0.6.1}/src/predikit/__init__.py +3 -2
  8. {predikit-0.5.2 → predikit-0.6.1}/src/predikit/cli.py +54 -0
  9. {predikit-0.5.2 → predikit-0.6.1}/src/predikit/coerce.py +4 -0
  10. predikit-0.6.1/src/predikit/mcp.py +115 -0
  11. {predikit-0.5.2 → predikit-0.6.1}/src/predikit/registry.py +19 -7
  12. predikit-0.6.1/tests/test_cli.py +219 -0
  13. {predikit-0.5.2 → predikit-0.6.1}/tests/test_coerce.py +11 -0
  14. predikit-0.6.1/tests/test_mcp.py +158 -0
  15. {predikit-0.5.2 → predikit-0.6.1}/tests/test_registry.py +15 -0
  16. {predikit-0.5.2 → predikit-0.6.1}/uv.lock +441 -21
  17. predikit-0.5.2/tests/test_cli.py +0 -98
  18. {predikit-0.5.2 → predikit-0.6.1}/.gitattributes +0 -0
  19. {predikit-0.5.2 → predikit-0.6.1}/.github/ISSUE_TEMPLATE/bug_report.md +0 -0
  20. {predikit-0.5.2 → predikit-0.6.1}/.github/ISSUE_TEMPLATE/feature_request.md +0 -0
  21. {predikit-0.5.2 → predikit-0.6.1}/.github/workflows/publish.yml +0 -0
  22. {predikit-0.5.2 → predikit-0.6.1}/.gitignore +0 -0
  23. {predikit-0.5.2 → predikit-0.6.1}/.pre-commit-config.yaml +0 -0
  24. {predikit-0.5.2 → predikit-0.6.1}/CLAUDE.md +0 -0
  25. {predikit-0.5.2 → predikit-0.6.1}/CONTRIBUTING.md +0 -0
  26. {predikit-0.5.2 → predikit-0.6.1}/LICENSE +0 -0
  27. {predikit-0.5.2 → predikit-0.6.1}/docs/logo.gif +0 -0
  28. {predikit-0.5.2 → predikit-0.6.1}/docs/logo.png +0 -0
  29. {predikit-0.5.2 → predikit-0.6.1}/examples/01_basic_sklearn.py +0 -0
  30. {predikit-0.5.2 → predikit-0.6.1}/examples/02_xgboost_regression.py +0 -0
  31. {predikit-0.5.2 → predikit-0.6.1}/examples/03_orlando_real_estate.py +0 -0
  32. {predikit-0.5.2 → predikit-0.6.1}/examples/04_confidence_routing.py +0 -0
  33. {predikit-0.5.2 → predikit-0.6.1}/examples/05_multi_model_ensemble.py +0 -0
  34. {predikit-0.5.2 → predikit-0.6.1}/examples/06_mlflow_loader.py +0 -0
  35. {predikit-0.5.2 → predikit-0.6.1}/examples/07_snowflake_loader.py +0 -0
  36. {predikit-0.5.2 → predikit-0.6.1}/src/predikit/ensemble.py +0 -0
  37. {predikit-0.5.2 → predikit-0.6.1}/src/predikit/exceptions.py +0 -0
  38. {predikit-0.5.2 → predikit-0.6.1}/src/predikit/exporters/__init__.py +0 -0
  39. {predikit-0.5.2 → predikit-0.6.1}/src/predikit/exporters/langchain.py +0 -0
  40. {predikit-0.5.2 → predikit-0.6.1}/src/predikit/exporters/openai.py +0 -0
  41. {predikit-0.5.2 → predikit-0.6.1}/src/predikit/introspect.py +0 -0
  42. {predikit-0.5.2 → predikit-0.6.1}/src/predikit/loaders/__init__.py +0 -0
  43. {predikit-0.5.2 → predikit-0.6.1}/src/predikit/loaders/mlflow.py +0 -0
  44. {predikit-0.5.2 → predikit-0.6.1}/src/predikit/loaders/snowflake.py +0 -0
  45. {predikit-0.5.2 → predikit-0.6.1}/src/predikit/tool.py +0 -0
  46. {predikit-0.5.2 → predikit-0.6.1}/tests/__init__.py +0 -0
  47. {predikit-0.5.2 → predikit-0.6.1}/tests/test_confidence.py +0 -0
  48. {predikit-0.5.2 → predikit-0.6.1}/tests/test_ensemble.py +0 -0
  49. {predikit-0.5.2 → predikit-0.6.1}/tests/test_exporters_langchain.py +0 -0
  50. {predikit-0.5.2 → predikit-0.6.1}/tests/test_exporters_openai.py +0 -0
  51. {predikit-0.5.2 → predikit-0.6.1}/tests/test_introspect.py +0 -0
  52. {predikit-0.5.2 → predikit-0.6.1}/tests/test_loaders_mlflow.py +0 -0
  53. {predikit-0.5.2 → predikit-0.6.1}/tests/test_loaders_snowflake.py +0 -0
  54. {predikit-0.5.2 → predikit-0.6.1}/tests/test_logging.py +0 -0
  55. {predikit-0.5.2 → predikit-0.6.1}/tests/test_tool.py +0 -0
  56. {predikit-0.5.2 → predikit-0.6.1}/tests/test_weighted_ensemble.py +0 -0
@@ -27,7 +27,19 @@
27
27
  "Bash(git stash *)",
28
28
  "Bash(./.venv/Scripts/python.exe -m pytest -q)",
29
29
  "Bash(./.venv/Scripts/python.exe -c \"import numpy; print\\(numpy.__version__\\); print\\(numpy.__file__\\)\")",
30
- "Bash(./.venv/Scripts/python.exe -m mypy --version)"
30
+ "Bash(./.venv/Scripts/python.exe -m mypy --version)",
31
+ "Bash(gh release *)",
32
+ "Bash(python -c \"import mcp; print\\('mcp', mcp.__version__ if hasattr\\(mcp,'__version__'\\) else 'installed'\\)\")",
33
+ "Bash(python -c \"import snowflake.ml; print\\('snowflake-ml installed'\\)\")",
34
+ "Bash(python -m pytest -q --cov=src/predikit --cov-report=term-missing)",
35
+ "Bash(py -0p)",
36
+ "Bash(where python *)",
37
+ "Bash(.venv/Scripts/python.exe -m pytest -q --cov=src/predikit --cov-report=term-missing)",
38
+ "Bash(command -v uv)",
39
+ "Bash(.venv/Scripts/python.exe -m pip download 'mcp>=1.28,<2' -d __TRACKED_VAR__/mcpdl --no-deps -q)",
40
+ "Bash(.venv/Scripts/python.exe -m pip show twine)",
41
+ "Bash(.venv/Scripts/python.exe -m pip install -q -U twine)",
42
+ "Bash(.venv/Scripts/python.exe -m twine check dist/*0.6.1*)"
31
43
  ],
32
44
  "additionalDirectories": [
33
45
  "C:\\AI\\predikit"
@@ -40,7 +40,7 @@ jobs:
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
@@ -6,6 +6,29 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
6
6
 
7
7
  ## [Unreleased]
8
8
 
9
+ ## [0.6.1] - 2026-08-15
10
+
11
+ ### Fixed
12
+ - 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.
13
+ - 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.
14
+ - `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`.
15
+ - `coerce_value()` converts `True` / `False` to `1` / `0` for `int` fields. `isinstance(True, int)` is true in Python, so bools previously passed through unconverted.
16
+
17
+ ### Added
18
+ - `--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`.
19
+ - `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.
20
+ - Python 3.13 in the CI matrix and the package classifiers.
21
+
22
+ ### Changed
23
+ - `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.
24
+ - Test coverage for `predikit serve`, which previously had none.
25
+
26
+ ## [0.6.0] - 2026-08-03
27
+
28
+ ### Added
29
+ - MCP server integration via `create_mcp_server()` and `predikit serve`.
30
+ - Optional `mcp` dependency for exposing `ToolRegistry` tools to MCP-compatible clients.
31
+
9
32
  ## [0.5.2] - 2026-08-01
10
33
 
11
34
  ### Fixed
@@ -1,6 +1,6 @@
1
- Metadata-Version: 2.4
1
+ Metadata-Version: 2.5
2
2
  Name: predikit
3
- Version: 0.5.2
3
+ Version: 0.6.1
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'
@@ -43,6 +45,8 @@ Requires-Dist: ruff>=0.4.0; extra == 'dev'
43
45
  Requires-Dist: xgboost>=1.7; extra == 'dev'
44
46
  Provides-Extra: langchain
45
47
  Requires-Dist: langchain-core>=0.1; extra == 'langchain'
48
+ Provides-Extra: mcp
49
+ Requires-Dist: mcp<2,>=1.28; extra == 'mcp'
46
50
  Provides-Extra: mlflow
47
51
  Requires-Dist: mlflow>=2.0; extra == 'mlflow'
48
52
  Provides-Extra: snowflake
@@ -324,6 +328,20 @@ response = client.chat.completions.create(
324
328
  agent = initialize_agent(tools=registry.to_langchain(), ...)
325
329
  ```
326
330
 
331
+ ### MCP server
332
+
333
+ Expose the same registry to MCP-compatible clients with the optional MCP extra:
334
+
335
+ ```bash
336
+ pip install predikit[mcp]
337
+ predikit serve my_tools:registry
338
+ ```
339
+
340
+ `my_tools:registry` must point to a `ToolRegistry`, or to a zero-argument function
341
+ that returns one. The default transport is `stdio`; use `--transport streamable-http`
342
+ when serving over HTTP. MCP tools use the same Pydantic input validation and model
343
+ execution as direct `invoke()` calls.
344
+
327
345
  ### Bool inputs from an LLM
328
346
 
329
347
  LLMs sometimes return `"yes"`, `"true"`, or `"1"` for boolean fields. predikit coerces these automatically before Pydantic validation:
@@ -463,7 +481,7 @@ See [`examples/03_orlando_real_estate.py`](examples/03_orlando_real_estate.py) f
463
481
  - [ ] HuggingFace / PyTorch / TensorFlow model support
464
482
  - [ ] Streaming inference support
465
483
  - [ ] OpenAI Assistants API integration
466
- - [ ] MCP server mode
484
+ - [x] MCP server mode (stdio and Streamable HTTP)
467
485
 
468
486
  ---
469
487
 
@@ -271,6 +271,20 @@ response = client.chat.completions.create(
271
271
  agent = initialize_agent(tools=registry.to_langchain(), ...)
272
272
  ```
273
273
 
274
+ ### MCP server
275
+
276
+ Expose the same registry to MCP-compatible clients with the optional MCP extra:
277
+
278
+ ```bash
279
+ pip install predikit[mcp]
280
+ predikit serve my_tools:registry
281
+ ```
282
+
283
+ `my_tools:registry` must point to a `ToolRegistry`, or to a zero-argument function
284
+ that returns one. The default transport is `stdio`; use `--transport streamable-http`
285
+ when serving over HTTP. MCP tools use the same Pydantic input validation and model
286
+ execution as direct `invoke()` calls.
287
+
274
288
  ### Bool inputs from an LLM
275
289
 
276
290
  LLMs sometimes return `"yes"`, `"true"`, or `"1"` for boolean fields. predikit coerces these automatically before Pydantic validation:
@@ -410,7 +424,7 @@ See [`examples/03_orlando_real_estate.py`](examples/03_orlando_real_estate.py) f
410
424
  - [ ] HuggingFace / PyTorch / TensorFlow model support
411
425
  - [ ] Streaming inference support
412
426
  - [ ] OpenAI Assistants API integration
413
- - [ ] MCP server mode
427
+ - [x] MCP server mode (stdio and Streamable HTTP)
414
428
 
415
429
  ---
416
430
 
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "predikit"
7
- version = "0.5.2"
7
+ version = "0.6.1"
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",
@@ -39,6 +40,7 @@ dependencies = [
39
40
 
40
41
  [project.optional-dependencies]
41
42
  langchain = ["langchain-core>=0.1"]
43
+ mcp = ["mcp>=1.28,<2"]
42
44
  xgboost = ["xgboost>=1.7"]
43
45
  mlflow = ["mlflow>=2.0"]
44
46
  snowflake = ["snowflake-ml-python>=1.0"]
@@ -49,6 +51,7 @@ dev = [
49
51
  "xgboost>=1.7",
50
52
  "langchain-core>=0.1",
51
53
  "pandas>=1.5",
54
+ "mcp>=1.28,<2",
52
55
  "mlflow>=2.0",
53
56
  "click>=8.0",
54
57
  "joblib>=1.2",
@@ -1,7 +1,8 @@
1
1
  from .ensemble import ModelEnsemble
2
2
  from .exceptions import LowConfidenceError
3
+ from .mcp import create_mcp_server
3
4
  from .registry import ToolRegistry
4
5
  from .tool import ModelTool
5
6
 
6
- __all__ = ["ModelTool", "ToolRegistry", "ModelEnsemble", "LowConfidenceError"]
7
- __version__ = "0.5.2"
7
+ __all__ = ["ModelTool", "ToolRegistry", "ModelEnsemble", "LowConfidenceError", "create_mcp_server"]
8
+ __version__ = "0.6.1"
@@ -1,5 +1,6 @@
1
1
  from __future__ import annotations
2
2
 
3
+ import importlib
3
4
  import json
4
5
  import sys
5
6
  from typing import Any
@@ -74,6 +75,59 @@ if _CLICK_AVAILABLE:
74
75
  else:
75
76
  click.echo("\nOpenAI schema: unavailable (fit model with a named DataFrame to enable)")
76
77
 
78
+ @cli.command()
79
+ @click.argument("registry_target")
80
+ @click.option("--name", default="predikit", show_default=True, help="MCP server name.")
81
+ @click.option(
82
+ "--transport",
83
+ type=click.Choice(["stdio", "streamable-http"]),
84
+ default="stdio",
85
+ show_default=True,
86
+ help="MCP transport to run.",
87
+ )
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:
101
+ """Serve a ToolRegistry from MODULE:ATTRIBUTE over MCP.
102
+
103
+ ATTRIBUTE may be a ToolRegistry instance or a zero-argument factory.
104
+ """
105
+ if ":" not in registry_target:
106
+ raise click.ClickException("REGISTRY_TARGET must use the form MODULE:ATTRIBUTE")
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.
112
+ try:
113
+ module = importlib.import_module(module_name)
114
+ except ImportError as err:
115
+ raise click.ClickException(f"Could not import module '{module_name}': {err}") from err
116
+ try:
117
+ registry = getattr(module, attribute_name)
118
+ except AttributeError as err:
119
+ raise click.ClickException(
120
+ f"Module '{module_name}' has no attribute '{attribute_name}'"
121
+ ) from err
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
+
77
131
  else:
78
132
 
79
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
 
@@ -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()]
@@ -0,0 +1,219 @@
1
+ import json
2
+
3
+ import pytest
4
+
5
+ joblib = pytest.importorskip("joblib")
6
+ click_testing = pytest.importorskip("click.testing")
7
+
8
+ from click.testing import CliRunner # noqa: E402
9
+ from sklearn.datasets import load_iris # noqa: E402
10
+ from sklearn.linear_model import LinearRegression, LogisticRegression # noqa: E402
11
+
12
+ from predikit.cli import cli # noqa: E402
13
+
14
+
15
+ @pytest.fixture
16
+ def runner():
17
+ return CliRunner()
18
+
19
+
20
+ @pytest.fixture
21
+ def reg_pkl(tmp_path):
22
+ """Regression model fitted with named features, so it has feature_names_in_."""
23
+ X, y = load_iris(return_X_y=True, as_frame=True)
24
+ model = LinearRegression().fit(X, y.astype(float))
25
+ path = tmp_path / "reg.pkl"
26
+ joblib.dump(model, path)
27
+ return str(path)
28
+
29
+
30
+ @pytest.fixture
31
+ def clf_pkl(tmp_path):
32
+ """Classifier fitted with named features."""
33
+ X, y = load_iris(return_X_y=True, as_frame=True)
34
+ model = LogisticRegression(max_iter=200).fit(X, y)
35
+ path = tmp_path / "clf.pkl"
36
+ joblib.dump(model, path)
37
+ return str(path)
38
+
39
+
40
+ @pytest.fixture
41
+ def unnamed_pkl(tmp_path):
42
+ """Model fitted without named features (numpy array)."""
43
+ X, y = load_iris(return_X_y=True)
44
+ model = LinearRegression().fit(X, y.astype(float))
45
+ path = tmp_path / "unnamed.pkl"
46
+ joblib.dump(model, path)
47
+ return str(path)
48
+
49
+
50
+ def test_inspect_regression_shows_model_type(runner, reg_pkl):
51
+ result = runner.invoke(cli, ["inspect", reg_pkl])
52
+ assert result.exit_code == 0, result.output
53
+ assert "LinearRegression" in result.output
54
+ assert "regression" in result.output
55
+
56
+
57
+ def test_inspect_classifier_shows_classes(runner, clf_pkl):
58
+ result = runner.invoke(cli, ["inspect", clf_pkl])
59
+ assert result.exit_code == 0, result.output
60
+ assert "classification" in result.output
61
+ assert "Classes" in result.output
62
+
63
+
64
+ def test_inspect_shows_feature_names(runner, reg_pkl):
65
+ result = runner.invoke(cli, ["inspect", reg_pkl])
66
+ assert "sepal length (cm)" in result.output
67
+
68
+
69
+ def test_inspect_emits_valid_openai_schema(runner, reg_pkl):
70
+ result = runner.invoke(cli, ["inspect", reg_pkl])
71
+ assert result.exit_code == 0
72
+ # Extract JSON from output (everything after "OpenAI schema:\n")
73
+ json_part = result.output.split("OpenAI schema:\n", 1)[1].strip()
74
+ schema = json.loads(json_part)
75
+ assert schema["type"] == "function"
76
+ assert "parameters" in schema["function"]
77
+
78
+
79
+ def test_inspect_custom_name_and_description(runner, reg_pkl):
80
+ result = runner.invoke(
81
+ cli, ["inspect", reg_pkl, "--name", "my_tool", "--description", "my desc"]
82
+ )
83
+ assert result.exit_code == 0
84
+ json_part = result.output.split("OpenAI schema:\n", 1)[1].strip()
85
+ schema = json.loads(json_part)
86
+ assert schema["function"]["name"] == "my_tool"
87
+ assert schema["function"]["description"] == "my desc"
88
+
89
+
90
+ def test_inspect_unnamed_model_skips_schema(runner, unnamed_pkl):
91
+ result = runner.invoke(cli, ["inspect", unnamed_pkl])
92
+ assert result.exit_code == 0
93
+ assert "unavailable" in result.output
94
+
95
+
96
+ def test_inspect_missing_file(runner):
97
+ result = runner.invoke(cli, ["inspect", "does_not_exist.pkl"])
98
+ assert result.exit_code != 0
99
+
100
+
101
+ def test_serve_requires_module_target(runner):
102
+ result = runner.invoke(cli, ["serve", "not-a-module-target"])
103
+ assert result.exit_code != 0
104
+ assert "MODULE:ATTRIBUTE" in result.output
105
+
106
+
107
+ @pytest.fixture
108
+ def registry_module(monkeypatch):
109
+ """A module exposing both a ToolRegistry instance and a factory returning one."""
110
+ import sys
111
+ import types
112
+
113
+ from pydantic import create_model
114
+
115
+ from predikit import ModelTool, ToolRegistry
116
+
117
+ class EchoModel:
118
+ def predict(self, X):
119
+ return [float(X[0][0])]
120
+
121
+ def build():
122
+ tool = ModelTool(
123
+ model=EchoModel(),
124
+ name="echo_value",
125
+ description="Echo a numeric value.",
126
+ input_schema=create_model("Input", value=(float, ...)),
127
+ output_name="value",
128
+ output_description="The value.",
129
+ )
130
+ return ToolRegistry([tool])
131
+
132
+ module = types.ModuleType("predikit_serve_fixture")
133
+ module.registry = build()
134
+ module.make_registry = build
135
+ monkeypatch.setitem(sys.modules, "predikit_serve_fixture", module)
136
+ return module
137
+
138
+
139
+ @pytest.fixture
140
+ def captured_run(monkeypatch):
141
+ """Stop short of actually serving; record how the server was configured."""
142
+ calls = {}
143
+
144
+ def fake_run(self, transport):
145
+ calls["transport"] = transport
146
+ calls["host"] = self.settings.host
147
+ calls["port"] = self.settings.port
148
+ calls["tools"] = sorted(self._tool_manager._tools)
149
+
150
+ pytest.importorskip("mcp")
151
+ from mcp.server.fastmcp import FastMCP
152
+
153
+ monkeypatch.setattr(FastMCP, "run", fake_run)
154
+ return calls
155
+
156
+
157
+ def test_serve_loads_a_registry_instance(runner, registry_module, captured_run):
158
+ result = runner.invoke(cli, ["serve", "predikit_serve_fixture:registry"])
159
+ assert result.exit_code == 0, result.output
160
+ assert captured_run["transport"] == "stdio"
161
+ assert captured_run["tools"] == ["echo_value"]
162
+
163
+
164
+ def test_serve_calls_a_registry_factory(runner, registry_module, captured_run):
165
+ result = runner.invoke(cli, ["serve", "predikit_serve_fixture:make_registry"])
166
+ assert result.exit_code == 0, result.output
167
+ assert captured_run["tools"] == ["echo_value"]
168
+
169
+
170
+ def test_serve_passes_host_and_port(runner, registry_module, captured_run):
171
+ result = runner.invoke(
172
+ cli,
173
+ [
174
+ "serve",
175
+ "predikit_serve_fixture:registry",
176
+ "--transport",
177
+ "streamable-http",
178
+ "--host",
179
+ "0.0.0.0",
180
+ "--port",
181
+ "9123",
182
+ ],
183
+ )
184
+ assert result.exit_code == 0, result.output
185
+ assert (captured_run["transport"], captured_run["host"], captured_run["port"]) == (
186
+ "streamable-http",
187
+ "0.0.0.0",
188
+ 9123,
189
+ )
190
+
191
+
192
+ def test_serve_reports_an_unimportable_module(runner):
193
+ result = runner.invoke(cli, ["serve", "no_such_module_xyz:registry"])
194
+ assert result.exit_code != 0
195
+ assert "Could not import module 'no_such_module_xyz'" in result.output
196
+
197
+
198
+ def test_serve_reports_a_missing_attribute(runner, registry_module):
199
+ result = runner.invoke(cli, ["serve", "predikit_serve_fixture:nope"])
200
+ assert result.exit_code != 0
201
+ assert "has no attribute 'nope'" in result.output
202
+
203
+
204
+ def test_serve_does_not_disguise_errors_from_the_registry_factory(runner, monkeypatch):
205
+ """An ImportError raised *inside* the target must not read as a bad target."""
206
+ import sys
207
+ import types
208
+
209
+ def broken():
210
+ raise ImportError("optional backend missing")
211
+
212
+ module = types.ModuleType("predikit_broken_fixture")
213
+ module.make_registry = broken
214
+ monkeypatch.setitem(sys.modules, "predikit_broken_fixture", module)
215
+
216
+ result = runner.invoke(cli, ["serve", "predikit_broken_fixture:make_registry"])
217
+ assert result.exit_code != 0
218
+ assert isinstance(result.exception, ImportError)
219
+ assert "Could not import module" not in result.output
@@ -55,3 +55,14 @@ def test_coerce_inputs_missing_feature_raises():
55
55
  meta = {"feature_names": ["x", "y"]}
56
56
  with pytest.raises(ValueError, match="missing"):
57
57
  coerce_inputs(validated, meta)
58
+
59
+
60
+ def test_bool_to_int_is_converted():
61
+ # isinstance(True, int) is True, so bools must be special-cased.
62
+ assert coerce_value(True, int) == 1
63
+ assert coerce_value(False, int) == 0
64
+ assert not isinstance(coerce_value(True, int), bool)
65
+
66
+
67
+ def test_bool_to_bool_is_unchanged():
68
+ assert coerce_value(True, bool) is True