predikit 0.4.4__tar.gz → 0.4.6__tar.gz

{predikit-0.4.4 → predikit-0.4.6}/.gitignore

@@ -1,13 +1,14 @@
 __pycache__/
 *.py[cod]
 *.egg-info/
-dist/
-build/
-.coverage
-coverage.xml
+dist/
+build/
+mlruns/
+.coverage
+coverage.xml
 htmlcov/
 .pytest_cache/
 .venv/
 venv/
 *.egg
-file.py
+file.py

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

@@ -4,11 +4,33 @@ 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.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
-## [Unreleased]
-
-## [0.4.4] - 2026-06-24
-
-### Changed
+## [Unreleased]
+
+## [0.4.6] - 2026-07-01
+
+### Fixed
+- Improved `ModelTool` input coercion for optional scalar annotations like `float | None`, `bool | None`, `int | None`, and `str | None`
+- Added test coverage for optional scalar string inputs to keep LLM-friendly coercion consistent across modern Pydantic schemas
+
+### Changed
+- Removed generated MLflow run artifacts from source control and ignored future `mlruns/` outputs
+- Cleaned package metadata and CLI text to avoid Unicode dash rendering issues in terminals and package indexes
+
+## [0.4.5] - 2026-07-01
+
+### Fixed
+- Added validation for `confidence_threshold` so values outside `0.0` to `1.0` fail at construction time
+- Added `ModelEnsemble` validation for duplicate `collect` outputs, mismatched aggregate output names, negative weights, and zero-total weighted strategies
+- Added `ToolRegistry` duplicate-name validation across tools and ensembles
+- Updated `ToolRegistry.get()` to retrieve both tools and ensembles with clearer missing-name errors
+- Cleaned up example lint issues and corrected stale `modelbridge[xgboost]` install text
+
+### Changed
+- Marked MLflow loader tests as opt-in integration tests so the default pytest suite stays fast and deterministic
+
+## [0.4.4] - 2026-06-24
+
+### Changed
 - README overhauled: added "Why predikit?" comparison table, "Works with" ecosystem section, GitHub star/fork badges, and a shipped-vs-planned roadmap
 - Quick start section replaces "30-second example" heading; `---` dividers added between major sections for better scanability
 - PyPI package description updated to a more descriptive one-liner

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

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: predikit
-Version: 0.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.
+Version: 0.4.6
+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
 Project-URL: Documentation, https://github.com/Tejas-TA/predikit#readme

{predikit-0.4.4 → predikit-0.4.6}/examples/01_basic_sklearn.py

@@ -1,4 +1,5 @@
 """Example 01 — Basic sklearn classifier wrapped as an OpenAI tool."""
+
 import json
 
 from pydantic import BaseModel, Field
@@ -36,7 +37,9 @@ print(json.dumps(tool.to_openai(), indent=2))
 
 # 5. Direct invocation
 print("\n=== Direct Invocation ===")
-result = tool.invoke({"sepal_length": 5.1, "sepal_width": 3.5, "petal_length": 1.4, "petal_width": 0.2})
+result = tool.invoke(
+    {"sepal_length": 5.1, "sepal_width": 3.5, "petal_length": 1.4, "petal_width": 0.2}
+)
 print(f"Result: {result}")
 
 # 6. Plain callable

{predikit-0.4.4 → predikit-0.4.6}/examples/02_xgboost_regression.py

@@ -1,7 +1,8 @@
 """
 Example 02 — XGBoost regression model wrapped as a tool.
-Requires: pip install modelbridge[xgboost]
+Requires: pip install predikit[xgboost]
 """
+
 import json
 
 from pydantic import BaseModel, Field
@@ -10,8 +11,8 @@ from sklearn.model_selection import train_test_split
 
 try:
     from xgboost import XGBRegressor
-except ImportError:
-    raise SystemExit("XGBoost not installed. Run: pip install modelbridge[xgboost]")
+except ImportError as err:
+    raise SystemExit("XGBoost not installed. Run: pip install predikit[xgboost]") from err
 
 from predikit import ModelTool
 
@@ -44,11 +45,13 @@ print("=== OpenAI Tool Schema ===")
 print(json.dumps(tool.to_openai(), indent=2))
 
 print("\n=== Invoke ===")
-result = tool.invoke({
-    "feature_1": float(X_test[0, 0]),
-    "feature_2": float(X_test[0, 1]),
-    "feature_3": float(X_test[0, 2]),
-    "feature_4": float(X_test[0, 3]),
-})
+result = tool.invoke(
+    {
+        "feature_1": float(X_test[0, 0]),
+        "feature_2": float(X_test[0, 1]),
+        "feature_3": float(X_test[0, 2]),
+        "feature_4": float(X_test[0, 3]),
+    }
+)
 print(f"Prediction: {result}")
 print(f"Actual:     {y_test[0]:.2f}")

{predikit-0.4.4 → predikit-0.4.6}/examples/03_orlando_real_estate.py

@@ -7,6 +7,7 @@ schema generation → registry export → direct invocation → callable.
 
 Requires: pip install xgboost
 """
+
 import json
 
 import numpy as np
@@ -16,8 +17,8 @@ from sklearn.model_selection import train_test_split
 
 try:
     from xgboost import XGBRegressor
-except ImportError:
-    raise SystemExit("XGBoost not installed. Run: pip install xgboost")
+except ImportError as err:
+    raise SystemExit("XGBoost not installed. Run: pip install xgboost") from err
 
 from predikit import ModelTool, ToolRegistry
 
@@ -27,31 +28,33 @@ from predikit import ModelTool, ToolRegistry
 rng = np.random.default_rng(42)
 n = 1_000
 
-sqft           = rng.integers(800, 4_000, n).astype(float)
-bedrooms       = rng.integers(1, 6, n).astype(float)
-bathrooms      = rng.choice([1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0], n)
-year_built     = rng.integers(1960, 2024, n).astype(float)
-has_pool       = rng.integers(0, 2, n).astype(float)
+sqft = rng.integers(800, 4_000, n).astype(float)
+bedrooms = rng.integers(1, 6, n).astype(float)
+bathrooms = rng.choice([1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0], n)
+year_built = rng.integers(1960, 2024, n).astype(float)
+has_pool = rng.integers(0, 2, n).astype(float)
 zip_code_group = rng.integers(0, 5, n).astype(float)  # 0=rural … 4=prime
 
 price = (
     80 * sqft
     + 15_000 * bedrooms
     + 10_000 * bathrooms
-    + 500   * (year_built - 1960)
+    + 500 * (year_built - 1960)
     + 20_000 * has_pool
     + 30_000 * zip_code_group
     + rng.normal(0, 15_000, n)
 ).clip(50_000, 1_500_000)
 
-df = pd.DataFrame({
-    "sqft":           sqft,
-    "bedrooms":       bedrooms,
-    "bathrooms":      bathrooms,
-    "year_built":     year_built,
-    "has_pool":       has_pool,
-    "zip_code_group": zip_code_group,
-})
+df = pd.DataFrame(
+    {
+        "sqft": sqft,
+        "bedrooms": bedrooms,
+        "bathrooms": bathrooms,
+        "year_built": year_built,
+        "has_pool": has_pool,
+        "zip_code_group": zip_code_group,
+    }
+)
 
 X_train, X_test, y_train, y_test = train_test_split(df, price, test_size=0.2, random_state=42)
 
@@ -61,17 +64,19 @@ X_train, X_test, y_train, y_test = train_test_split(df, price, test_size=0.2, ra
 model = XGBRegressor(n_estimators=200, learning_rate=0.05, random_state=42)
 model.fit(X_train, y_train)
 
+
 # ---------------------------------------------------------------------------
 # 3. Input schema — field names must match DataFrame columns exactly
 # ---------------------------------------------------------------------------
 class OrlandoHouseInput(BaseModel):
-    sqft:           float = Field(description="Total square footage of the home")
-    bedrooms:       float = Field(description="Number of bedrooms")
-    bathrooms:      float = Field(description="Number of bathrooms (0.5 increments)")
-    year_built:     float = Field(description="Year the home was built")
-    has_pool:       float = Field(description="1 if the home has a pool, else 0")
+    sqft: float = Field(description="Total square footage of the home")
+    bedrooms: float = Field(description="Number of bedrooms")
+    bathrooms: float = Field(description="Number of bathrooms (0.5 increments)")
+    year_built: float = Field(description="Year the home was built")
+    has_pool: float = Field(description="1 if the home has a pool, else 0")
     zip_code_group: float = Field(description="Area cluster 0–4 (0=rural, 4=prime location)")
 
+
 # ---------------------------------------------------------------------------
 # 4. Wrap as ModelTool and register
 # ---------------------------------------------------------------------------
@@ -99,11 +104,11 @@ print(json.dumps(registry.to_openai(), indent=2))
 # 6. Run a prediction
 # ---------------------------------------------------------------------------
 sample = {
-    "sqft":           2_200.0,
-    "bedrooms":       3.0,
-    "bathrooms":      2.0,
-    "year_built":     2005.0,
-    "has_pool":       1.0,
+    "sqft": 2_200.0,
+    "bedrooms": 3.0,
+    "bathrooms": 2.0,
+    "year_built": 2005.0,
+    "has_pool": 1.0,
     "zip_code_group": 3.0,
 }
 

{predikit-0.4.4 → predikit-0.4.6}/examples/04_confidence_routing.py

@@ -9,6 +9,7 @@ Demonstrates the three on_low_confidence modes:
 threshold=0.999 guarantees the low-confidence path fires for illustration.
 In production, set this between 0.75 and 0.95.
 """
+
 from pydantic import BaseModel, Field
 from sklearn.datasets import load_iris
 from sklearn.linear_model import LogisticRegression
@@ -26,9 +27,9 @@ fallback_clf = GaussianNB().fit(X, y)
 
 class IrisInput(BaseModel):
     sepal_length: float = Field(description="Sepal length in cm")
-    sepal_width:  float = Field(description="Sepal width in cm")
+    sepal_width: float = Field(description="Sepal width in cm")
     petal_length: float = Field(description="Petal length in cm")
-    petal_width:  float = Field(description="Petal width in cm")
+    petal_width: float = Field(description="Petal width in cm")
 
 
 fallback_tool = ModelTool(
@@ -61,7 +62,9 @@ print("=== Mode: warn ===")
 result = warn_tool.invoke(sample)
 print(result)
 if result.get("_low_confidence"):
-    print(f"  → confidence {result['_confidence']:.3f} below threshold; agent should seek clarification")
+    print(
+        f"  → confidence {result['_confidence']:.3f} below threshold; agent should seek clarification"
+    )
 
 # ---------------------------------------------------------------------------
 # Mode 2: raise — lets the calling agent handle uncertainty explicitly

{predikit-0.4.4 → predikit-0.4.6}/examples/05_multi_model_ensemble.py

@@ -8,6 +8,7 @@ Demonstrates ModelEnsemble with all three strategies:
 
 Requires: pip install predikit[xgboost]
 """
+
 import json
 
 import numpy as np
@@ -19,8 +20,8 @@ from sklearn.model_selection import train_test_split
 
 try:
     from xgboost import XGBRegressor
-except ImportError:
-    raise SystemExit("XGBoost not installed. Run: pip install predikit[xgboost]")
+except ImportError as err:
+    raise SystemExit("XGBoost not installed. Run: pip install predikit[xgboost]") from err
 
 from predikit import ModelEnsemble, ModelTool, ToolRegistry
 
@@ -29,37 +30,51 @@ from predikit import ModelEnsemble, ModelTool, ToolRegistry
 # ---------------------------------------------------------------------------
 rng = np.random.default_rng(42)
 n = 1_000
-sqft           = rng.integers(800, 4_000, n).astype(float)
-bedrooms       = rng.integers(1, 6, n).astype(float)
-bathrooms      = rng.choice([1.0, 1.5, 2.0, 2.5, 3.0], n)
-year_built     = rng.integers(1960, 2024, n).astype(float)
-has_pool       = rng.integers(0, 2, n).astype(float)
+sqft = rng.integers(800, 4_000, n).astype(float)
+bedrooms = rng.integers(1, 6, n).astype(float)
+bathrooms = rng.choice([1.0, 1.5, 2.0, 2.5, 3.0], n)
+year_built = rng.integers(1960, 2024, n).astype(float)
+has_pool = rng.integers(0, 2, n).astype(float)
 zip_code_group = rng.integers(0, 5, n).astype(float)
 price = (
-    80 * sqft + 15_000 * bedrooms + 10_000 * bathrooms
-    + 500 * (year_built - 1960) + 20_000 * has_pool
-    + 30_000 * zip_code_group + rng.normal(0, 15_000, n)
+    80 * sqft
+    + 15_000 * bedrooms
+    + 10_000 * bathrooms
+    + 500 * (year_built - 1960)
+    + 20_000 * has_pool
+    + 30_000 * zip_code_group
+    + rng.normal(0, 15_000, n)
 ).clip(50_000, 1_500_000)
 
-df = pd.DataFrame({
-    "sqft": sqft, "bedrooms": bedrooms, "bathrooms": bathrooms,
-    "year_built": year_built, "has_pool": has_pool, "zip_code_group": zip_code_group,
-})
+df = pd.DataFrame(
+    {
+        "sqft": sqft,
+        "bedrooms": bedrooms,
+        "bathrooms": bathrooms,
+        "year_built": year_built,
+        "has_pool": has_pool,
+        "zip_code_group": zip_code_group,
+    }
+)
 X_train, X_test, y_train, _ = train_test_split(df, price, test_size=0.2, random_state=42)
 
 
 class HouseInput(BaseModel):
-    sqft:           float = Field(description="Square footage")
-    bedrooms:       float = Field(description="Number of bedrooms")
-    bathrooms:      float = Field(description="Number of bathrooms")
-    year_built:     float = Field(description="Year built")
-    has_pool:       float = Field(description="1 if has pool, else 0")
+    sqft: float = Field(description="Square footage")
+    bedrooms: float = Field(description="Number of bedrooms")
+    bathrooms: float = Field(description="Number of bathrooms")
+    year_built: float = Field(description="Year built")
+    has_pool: float = Field(description="1 if has pool, else 0")
     zip_code_group: float = Field(description="Area cluster 0–4")
 
 
 sample = {
-    "sqft": 2200.0, "bedrooms": 3.0, "bathrooms": 2.0,
-    "year_built": 2005.0, "has_pool": 1.0, "zip_code_group": 3.0,
+    "sqft": 2200.0,
+    "bedrooms": 3.0,
+    "bathrooms": 2.0,
+    "year_built": 2005.0,
+    "has_pool": 1.0,
+    "zip_code_group": 3.0,
 }
 
 # ---------------------------------------------------------------------------
@@ -67,13 +82,19 @@ sample = {
 # ---------------------------------------------------------------------------
 xgb_tool = ModelTool(
     model=XGBRegressor(n_estimators=100, random_state=42).fit(X_train, y_train),
-    name="xgb_price", description="XGBoost price estimate",
-    input_schema=HouseInput, output_name="xgb_price_usd", output_description="XGBoost price",
+    name="xgb_price",
+    description="XGBoost price estimate",
+    input_schema=HouseInput,
+    output_name="xgb_price_usd",
+    output_description="XGBoost price",
 )
 lr_tool = ModelTool(
     model=LinearRegression().fit(X_train, y_train),
-    name="lr_price", description="Linear regression price estimate",
-    input_schema=HouseInput, output_name="lr_price_usd", output_description="LR price",
+    name="lr_price",
+    description="Linear regression price estimate",
+    input_schema=HouseInput,
+    output_name="lr_price_usd",
+    output_description="LR price",
 )
 
 collect_ensemble = ModelEnsemble(
@@ -95,8 +116,11 @@ print(f"  Linear:   ${result['lr_price_usd']:,.0f}")
 mean_tools = [
     ModelTool(
         model=XGBRegressor(n_estimators=100, random_state=seed).fit(X_train, y_train),
-        name=f"xgb_{seed}", description=f"XGBoost seed={seed}",
-        input_schema=HouseInput, output_name="price_usd", output_description="price",
+        name=f"xgb_{seed}",
+        description=f"XGBoost seed={seed}",
+        input_schema=HouseInput,
+        output_name="price_usd",
+        output_description="price",
     )
     for seed in [0, 1]
 ]
@@ -120,16 +144,19 @@ X_iris, y_iris = load_iris(return_X_y=True)
 
 class IrisInput(BaseModel):
     sepal_length: float
-    sepal_width:  float
+    sepal_width: float
     petal_length: float
-    petal_width:  float
+    petal_width: float
 
 
 clf_tools = [
     ModelTool(
         model=LogisticRegression(max_iter=200, C=c).fit(X_iris, y_iris),
-        name=f"clf_C{c}", description=f"LR C={c}",
-        input_schema=IrisInput, output_name="species", output_description="species",
+        name=f"clf_C{c}",
+        description=f"LR C={c}",
+        input_schema=IrisInput,
+        output_name="species",
+        output_description="species",
     )
     for c in [0.1, 1.0, 10.0]
 ]

{predikit-0.4.4 → predikit-0.4.6}/examples/06_mlflow_loader.py

@@ -11,17 +11,16 @@ Run:
     pip install predikit[mlflow]
     python examples/06_mlflow_loader.py
 """
+
 import json
 import tempfile
 
 import mlflow
 import mlflow.sklearn
-import numpy as np
 from pydantic import BaseModel, Field
 from sklearn.datasets import load_iris
 from sklearn.linear_model import LogisticRegression
 
-from predikit import LowConfidenceError
 from predikit.loaders import from_mlflow
 
 # ---------------------------------------------------------------------------
@@ -48,9 +47,9 @@ print(f"Logged model → {model_uri}\n")
 # ---------------------------------------------------------------------------
 class IrisInput(BaseModel):
     sepal_length: float = Field(description="Sepal length in cm")
-    sepal_width:  float = Field(description="Sepal width in cm")
+    sepal_width: float = Field(description="Sepal width in cm")
     petal_length: float = Field(description="Petal length in cm")
-    petal_width:  float = Field(description="Petal width in cm")
+    petal_width: float = Field(description="Petal width in cm")
 
 
 # ---------------------------------------------------------------------------

{predikit-0.4.4 → predikit-0.4.6}/examples/07_snowflake_loader.py

@@ -15,13 +15,14 @@ For a real Snowflake connection:
     #   from snowflake.snowpark import Session
     #   session = Session.builder.configs({...}).create()
 """
+
 import json
 from unittest.mock import MagicMock
 
 import numpy as np
 from pydantic import BaseModel, Field
 
-from predikit.loaders.snowflake import _SnowflakeShim, from_snowflake
+from predikit.loaders.snowflake import from_snowflake
 
 # ---------------------------------------------------------------------------
 # Mock a Snowflake session and model (stand-in for real Snowpark objects)
@@ -36,6 +37,7 @@ mock_sf_model.predict.return_value = np.array([0.73])
 def _patch_registry(mock_sf_model):
     """Context manager that replaces Registry with a stub returning mock_sf_model."""
     from unittest.mock import patch
+
     registry_mock = MagicMock()
     registry_mock.get_model.return_value.version.return_value = mock_sf_model
     return patch("predikit.loaders.snowflake.Registry", return_value=registry_mock)
@@ -45,9 +47,9 @@ def _patch_registry(mock_sf_model):
 # Define the input schema matching the Snowflake model's feature columns
 # ---------------------------------------------------------------------------
 class MemberInput(BaseModel):
-    tenure_months:  float = Field(description="Months as a Vacation Ownership member")
+    tenure_months: float = Field(description="Months as a Vacation Ownership member")
     trips_last_year: float = Field(description="Number of trips taken in the past 12 months")
-    avg_spend:      float = Field(description="Average spend per trip in USD")
+    avg_spend: float = Field(description="Average spend per trip in USD")
 
 
 # ---------------------------------------------------------------------------

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

@@ -4,8 +4,8 @@ build-backend = "hatchling.build"
 
 [project]
 name = "predikit"
-version = "0.4.4"
-description = "Wrap any trained sklearn/XGBoost model as an LLM-callable tool — auto-generated JSON schemas, Pydantic I/O, OpenAI & LangChain ready, zero boilerplate."
+version = "0.4.6"
+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"}
 requires-python = ">=3.10"
@@ -70,8 +70,12 @@ Changelog = "https://github.com/Tejas-TA/predikit/releases"
 [tool.hatch.build.targets.wheel]
 packages = ["src/predikit"]
 
-[tool.pytest.ini_options]
-testpaths = ["tests"]
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+addopts = ["-m", "not integration"]
+markers = [
+    "integration: tests that exercise slower external runtimes or registries",
+]
 
 [tool.ruff]
 line-length = 100

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

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

{predikit-0.4.4 → predikit-0.4.6}/src/predikit/cli.py

@@ -16,7 +16,7 @@ if _CLICK_AVAILABLE:
 
     @click.group()
     def cli() -> None:
-        """predikit — ML model utilities for LLM agents."""
+        """predikit: ML model utilities for LLM agents."""
 
     @cli.command()
     @click.argument("model_path", type=click.Path(exists=True))
@@ -53,7 +53,7 @@ if _CLICK_AVAILABLE:
             for fname in meta["feature_names"]:
                 click.echo(f"  {fname}")
         else:
-            click.echo("Feature names: (none — fit the model with a named DataFrame to enable)")
+            click.echo("Feature names: (none - fit the model with a named DataFrame to enable)")
 
         if meta["classes"] is not None:
             click.echo(f"Classes:  {meta['classes']}")

{predikit-0.4.4 → predikit-0.4.6}/src/predikit/ensemble.py

@@ -10,6 +10,7 @@ from .exporters.openai import to_openai_schema
 from .tool import ModelTool
 
 _VALID_STRATEGIES = {"collect", "mean", "weighted_mean", "vote", "weighted_vote"}
+_WEIGHTED_STRATEGIES = {"weighted_mean", "weighted_vote"}
 
 
 class ModelEnsemble:
@@ -33,6 +34,24 @@ class ModelEnsemble:
             raise ValueError(
                 f"weights length ({len(weights)}) must match number of tools ({len(tools)})."
             )
+        if weights is not None:
+            if any(w < 0 for w in weights):
+                raise ValueError("weights must be non-negative.")
+            if strategy in _WEIGHTED_STRATEGIES and sum(weights) <= 0:
+                raise ValueError("weights must sum to a positive value.")
+        output_names = [tool.output_name for tool in tools]
+        if strategy == "collect":
+            duplicates = sorted({name for name in output_names if output_names.count(name) > 1})
+            if duplicates:
+                raise ValueError(
+                    "collect strategy requires unique output_name values; "
+                    f"duplicates: {duplicates}."
+                )
+        elif len(set(output_names)) != 1:
+            raise ValueError(
+                f"{strategy} strategy requires all tools to share the same output_name; "
+                f"got {output_names}."
+            )
         self.tools = tools
         self.name = name
         self.description = description

{predikit-0.4.4 → predikit-0.4.6}/src/predikit/registry.py

@@ -3,6 +3,8 @@ from __future__ import annotations
 from .ensemble import ModelEnsemble
 from .tool import ModelTool
 
+RegistryItem = ModelTool | ModelEnsemble
+
 
 class ToolRegistry:
     """Bundles multiple ModelTools (and optional ModelEnsembles) for bulk export."""
@@ -12,13 +14,20 @@ class ToolRegistry:
         tools: list[ModelTool],
         ensembles: list[ModelEnsemble] | None = None,
     ) -> None:
+        names = [t.name for t in tools] + [e.name for e in (ensembles or [])]
+        duplicates = sorted({name for name in names if names.count(name) > 1})
+        if duplicates:
+            raise ValueError(f"Duplicate tool or ensemble names are not allowed: {duplicates}.")
         self._tools: dict[str, ModelTool] = {t.name: t for t in tools}
         self._ensembles: dict[str, ModelEnsemble] = {e.name: e for e in (ensembles or [])}
 
-    def get(self, name: str) -> ModelTool:
-        if name not in self._tools:
-            raise KeyError(f"No tool named '{name}'. Available: {list(self._tools)}")
-        return self._tools[name]
+    def get(self, name: str) -> RegistryItem:
+        if name in self._tools:
+            return self._tools[name]
+        if name in self._ensembles:
+            return self._ensembles[name]
+        available = list(self._tools) + list(self._ensembles)
+        raise KeyError(f"No tool or ensemble named '{name}'. Available: {available}")
 
     def to_openai(self) -> list[dict]:
         return [t.to_openai() for t in self._tools.values()] + [

{predikit-0.4.4 → predikit-0.4.6}/src/predikit/tool.py

@@ -4,7 +4,8 @@ import asyncio
 import time
 import warnings
 from collections.abc import Callable
-from typing import Any
+from types import UnionType
+from typing import Any, Union, get_args, get_origin
 
 import numpy as np
 from pydantic import BaseModel
@@ -17,6 +18,7 @@ from .introspect import introspect
 
 _VALID_ON_LOW_CONFIDENCE = {"warn", "raise", "fallback"}
 _SCALAR_TYPES = (bool, int, float, str)
+_NONE_TYPE = type(None)
 
 
 class ModelTool:
@@ -35,6 +37,10 @@ class ModelTool:
         fallback_tool: ModelTool | None = None,
         verbose: bool = False,
     ) -> None:
+        if confidence_threshold is not None and not 0.0 <= confidence_threshold <= 1.0:
+            raise ValueError(
+                f"confidence_threshold must be between 0.0 and 1.0, got {confidence_threshold!r}"
+            )
         if on_low_confidence not in _VALID_ON_LOW_CONFIDENCE:
             raise ValueError(
                 f"on_low_confidence must be one of {_VALID_ON_LOW_CONFIDENCE}, got {on_low_confidence!r}"
@@ -139,9 +145,10 @@ class ModelTool:
         for k, v in input_dict.items():
             field = fields.get(k)
             annotation = getattr(field, "annotation", None) if field else None
-            if annotation in _SCALAR_TYPES:
+            scalar_type = _scalar_type(annotation)
+            if scalar_type is not None and v is not None:
                 try:
-                    result[k] = coerce_value(v, annotation)
+                    result[k] = coerce_value(v, scalar_type)
                 except (ValueError, TypeError):
                     result[k] = v
             else:
@@ -158,3 +165,15 @@ class ModelTool:
             except ImportError:
                 pass
         return np.array([features])
+
+
+def _scalar_type(annotation: Any) -> type | None:
+    if annotation in _SCALAR_TYPES:
+        return annotation
+
+    if get_origin(annotation) in (Union, UnionType):
+        args = [arg for arg in get_args(annotation) if arg is not _NONE_TYPE]
+        if len(args) == 1:
+            return _scalar_type(args[0])
+
+    return None