predikit 0.1.0__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.
@@ -0,0 +1,27 @@
1
+ name: Publish to PyPI
2
+
3
+ on:
4
+ release:
5
+ types: [published]
6
+
7
+ jobs:
8
+ publish:
9
+ runs-on: ubuntu-latest
10
+ environment: pypi
11
+ permissions:
12
+ id-token: write
13
+
14
+ steps:
15
+ - uses: actions/checkout@v4
16
+
17
+ - uses: actions/setup-python@v5
18
+ with:
19
+ python-version: "3.12"
20
+
21
+ - name: Build
22
+ run: |
23
+ pip install hatchling
24
+ python -m hatchling build
25
+
26
+ - name: Publish to PyPI
27
+ uses: pypa/gh-action-pypi-publish@release/v1
@@ -0,0 +1,27 @@
1
+ name: Test
2
+
3
+ on:
4
+ push:
5
+ branches: [main]
6
+ pull_request:
7
+ branches: [main]
8
+
9
+ jobs:
10
+ test:
11
+ runs-on: ubuntu-latest
12
+ strategy:
13
+ matrix:
14
+ python-version: ["3.10", "3.11", "3.12"]
15
+
16
+ steps:
17
+ - uses: actions/checkout@v4
18
+
19
+ - uses: actions/setup-python@v5
20
+ with:
21
+ python-version: ${{ matrix.python-version }}
22
+
23
+ - name: Install dependencies
24
+ run: pip install -e ".[dev]"
25
+
26
+ - name: Run tests
27
+ run: pytest --cov=src/predikit --cov-report=term-missing
@@ -0,0 +1,12 @@
1
+ __pycache__/
2
+ *.py[cod]
3
+ *.egg-info/
4
+ dist/
5
+ build/
6
+ .coverage
7
+ coverage.xml
8
+ htmlcov/
9
+ .pytest_cache/
10
+ .venv/
11
+ venv/
12
+ *.egg
predikit-0.1.0/LICENSE ADDED
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Tejas Tumakuru Ashok
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,233 @@
1
+ Metadata-Version: 2.4
2
+ Name: predikit
3
+ Version: 0.1.0
4
+ Summary: Turn any trained sklearn/XGBoost model into an LLM-callable tool with auto-generated schemas and typed I/O.
5
+ Project-URL: Homepage, https://github.com/Tejas-TA/predikit
6
+ Project-URL: Repository, https://github.com/Tejas-TA/predikit
7
+ Author-email: Tejas Tumakuru Ashok <tejas.tumakuruashok@travelandleisure.com>
8
+ License: MIT
9
+ License-File: LICENSE
10
+ Keywords: agents,function-calling,llm,ml-tools,sklearn,xgboost
11
+ Classifier: Development Status :: 3 - Alpha
12
+ Classifier: Intended Audience :: Developers
13
+ Classifier: Intended Audience :: Science/Research
14
+ Classifier: License :: OSI Approved :: MIT License
15
+ Classifier: Programming Language :: Python :: 3
16
+ Classifier: Programming Language :: Python :: 3.10
17
+ Classifier: Programming Language :: Python :: 3.11
18
+ Classifier: Programming Language :: Python :: 3.12
19
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
20
+ Requires-Python: >=3.10
21
+ Requires-Dist: numpy>=1.24
22
+ Requires-Dist: pydantic>=2.0
23
+ Requires-Dist: scikit-learn>=1.2
24
+ Provides-Extra: dev
25
+ Requires-Dist: langchain-core>=0.1; extra == 'dev'
26
+ Requires-Dist: pandas>=1.5; extra == 'dev'
27
+ Requires-Dist: pytest-cov>=4.0; extra == 'dev'
28
+ Requires-Dist: pytest>=7.0; extra == 'dev'
29
+ Requires-Dist: xgboost>=1.7; extra == 'dev'
30
+ Provides-Extra: langchain
31
+ Requires-Dist: langchain-core>=0.1; extra == 'langchain'
32
+ Provides-Extra: xgboost
33
+ Requires-Dist: xgboost>=1.7; extra == 'xgboost'
34
+ Description-Content-Type: text/markdown
35
+
36
+ # predikit
37
+
38
+ Turn any trained scikit-learn or XGBoost model into an LLM-callable tool — auto-generated JSON schemas, typed I/O, zero boilerplate.
39
+
40
+ ```python
41
+ tool = ModelTool(model=clf, name="classify_iris", ...)
42
+ tool.to_openai() # OpenAI function schema, ready to pass to the API
43
+ tool.invoke({"sqft": 2200}) # → {"price_usd": 370730}
44
+ ```
45
+
46
+ ## Install
47
+
48
+ ```bash
49
+ pip install predikit
50
+
51
+ # With XGBoost support
52
+ pip install predikit[xgboost]
53
+
54
+ # With LangChain support
55
+ pip install predikit[langchain]
56
+ ```
57
+
58
+ ## 30-second example
59
+
60
+ ```python
61
+ from pydantic import BaseModel, Field
62
+ from sklearn.datasets import load_iris
63
+ from sklearn.linear_model import LogisticRegression
64
+ from predikit import ModelTool
65
+
66
+ # Train
67
+ X, y = load_iris(return_X_y=True)
68
+ clf = LogisticRegression(max_iter=200).fit(X, y)
69
+
70
+ # Define what the LLM will pass in
71
+ class IrisInput(BaseModel):
72
+ sepal_length: float = Field(description="Sepal length in cm")
73
+ sepal_width: float = Field(description="Sepal width in cm")
74
+ petal_length: float = Field(description="Petal length in cm")
75
+ petal_width: float = Field(description="Petal width in cm")
76
+
77
+ # Wrap the model
78
+ tool = ModelTool(
79
+ model=clf,
80
+ name="classify_iris",
81
+ description="Classify an iris flower: 0=setosa, 1=versicolor, 2=virginica.",
82
+ input_schema=IrisInput,
83
+ output_name="species",
84
+ output_description="Predicted species index",
85
+ )
86
+
87
+ # Get an OpenAI-ready schema
88
+ import json
89
+ print(json.dumps(tool.to_openai(), indent=2))
90
+
91
+ # Call it directly
92
+ tool.invoke({
93
+ "sepal_length": 5.1, "sepal_width": 3.5,
94
+ "petal_length": 1.4, "petal_width": 0.2,
95
+ })
96
+ # → {"species": 0}
97
+ ```
98
+
99
+ ## Core API
100
+
101
+ ### `ModelTool`
102
+
103
+ ```python
104
+ ModelTool(
105
+ model, # fitted sklearn-compatible estimator
106
+ name: str, # tool name the LLM sees
107
+ description: str, # tool description the LLM sees
108
+ input_schema, # Pydantic BaseModel describing inputs
109
+ output_name: str, # key for the prediction in the returned dict
110
+ output_description: str,
111
+ )
112
+ ```
113
+
114
+ | Method | Returns | What it does |
115
+ |--------|---------|--------------|
116
+ | `.invoke(input_dict)` | `dict` | Validates → predicts → returns `{output_name: value}` |
117
+ | `.to_openai()` | `dict` | OpenAI function-calling schema |
118
+ | `.to_langchain()` | `StructuredTool` | LangChain tool |
119
+ | `.to_callable()` | `Callable` | Plain Python function |
120
+
121
+ ### `ToolRegistry`
122
+
123
+ Group multiple tools for bulk export:
124
+
125
+ ```python
126
+ registry = ToolRegistry([price_tool, risk_tool])
127
+ registry.to_openai() # → list[dict], pass directly to OpenAI
128
+ registry.to_langchain() # → list[StructuredTool]
129
+ registry.get("name") # → ModelTool
130
+ ```
131
+
132
+ ## Field naming rule
133
+
134
+ **Your Pydantic schema field names must exactly match the column names the model was trained on.**
135
+
136
+ predikit maps inputs to features by name, not position. If you trained on a DataFrame with columns `["sqft", "bedrooms"]`, your schema fields must be `sqft` and `bedrooms` — not `sq_ft`, not `Sqft`.
137
+
138
+ ```python
139
+ # ✓ Columns match: sqft, bedrooms, bathrooms
140
+ class GoodInput(BaseModel):
141
+ sqft: float
142
+ bedrooms: float
143
+ bathrooms: float
144
+
145
+ # ✗ Name mismatch — raises ValueError at runtime
146
+ class BadInput(BaseModel):
147
+ square_footage: float # model expects "sqft"
148
+ beds: float # model expects "bedrooms"
149
+ baths: float # model expects "bathrooms"
150
+ ```
151
+
152
+ When there's a mismatch, predikit tells you exactly which names are wrong:
153
+
154
+ ```
155
+ ValueError: Input schema is missing model features: ['sqft', 'bedrooms'].
156
+ Schema has: ['square_footage', 'beds', 'bathrooms'], model expects: ['sqft', 'bedrooms', 'bathrooms']
157
+ ```
158
+
159
+ > **Tip:** If you trained with a numpy array (no DataFrame), predikit has no feature names to check — it uses your schema's field definition order instead.
160
+
161
+ ## Cookbook
162
+
163
+ ### XGBoost regression
164
+
165
+ ```python
166
+ from xgboost import XGBRegressor
167
+ from predikit import ModelTool
168
+
169
+ reg = XGBRegressor().fit(X_train, y_train)
170
+
171
+ class HouseInput(BaseModel):
172
+ sqft: float
173
+ bedrooms: float
174
+ year_built: float
175
+
176
+ tool = ModelTool(
177
+ model=reg,
178
+ name="price_estimate",
179
+ description="Predict home price in USD.",
180
+ input_schema=HouseInput,
181
+ output_name="price_usd",
182
+ output_description="Predicted sale price in USD",
183
+ )
184
+ ```
185
+
186
+ ### Multiple tools in one registry
187
+
188
+ ```python
189
+ registry = ToolRegistry([price_tool, risk_tool, demand_tool])
190
+
191
+ # OpenAI
192
+ response = client.chat.completions.create(
193
+ model="gpt-4o",
194
+ tools=registry.to_openai(),
195
+ ...
196
+ )
197
+
198
+ # LangChain
199
+ agent = initialize_agent(tools=registry.to_langchain(), ...)
200
+ ```
201
+
202
+ ### Bool inputs from an LLM
203
+
204
+ LLMs sometimes return `"yes"`, `"true"`, or `"1"` for boolean fields. predikit coerces these automatically before Pydantic validation:
205
+
206
+ ```python
207
+ class Input(BaseModel):
208
+ has_pool: bool
209
+
210
+ tool.invoke({"has_pool": "yes"}) # → coerced to True
211
+ tool.invoke({"has_pool": "false"}) # → coerced to False
212
+ tool.invoke({"has_pool": "maybe"}) # → raises ValueError with clear message
213
+ ```
214
+
215
+ Supported strings: `true/false`, `yes/no`, `1/0`, `on/off`.
216
+
217
+ ### Orlando real estate demo
218
+
219
+ See [`examples/03_orlando_real_estate.py`](examples/03_orlando_real_estate.py) for a full end-to-end walkthrough: synthetic dataset → XGBoost training → `ModelTool` → registry → OpenAI schema → prediction.
220
+
221
+ ## Roadmap
222
+
223
+ Intentionally out of scope for v0.1 — planned for later releases:
224
+
225
+ - Confidence-aware routing & fallback
226
+ - Multi-model synthesis (agent calls several, reconciles results)
227
+ - MLflow / Snowflake Model Registry integration
228
+ - HuggingFace / PyTorch / TensorFlow support
229
+ - Async invocation
230
+
231
+ ## License
232
+
233
+ MIT © Tejas Tumakuru Ashok
@@ -0,0 +1,198 @@
1
+ # predikit
2
+
3
+ Turn any trained scikit-learn or XGBoost model into an LLM-callable tool — auto-generated JSON schemas, typed I/O, zero boilerplate.
4
+
5
+ ```python
6
+ tool = ModelTool(model=clf, name="classify_iris", ...)
7
+ tool.to_openai() # OpenAI function schema, ready to pass to the API
8
+ tool.invoke({"sqft": 2200}) # → {"price_usd": 370730}
9
+ ```
10
+
11
+ ## Install
12
+
13
+ ```bash
14
+ pip install predikit
15
+
16
+ # With XGBoost support
17
+ pip install predikit[xgboost]
18
+
19
+ # With LangChain support
20
+ pip install predikit[langchain]
21
+ ```
22
+
23
+ ## 30-second example
24
+
25
+ ```python
26
+ from pydantic import BaseModel, Field
27
+ from sklearn.datasets import load_iris
28
+ from sklearn.linear_model import LogisticRegression
29
+ from predikit import ModelTool
30
+
31
+ # Train
32
+ X, y = load_iris(return_X_y=True)
33
+ clf = LogisticRegression(max_iter=200).fit(X, y)
34
+
35
+ # Define what the LLM will pass in
36
+ class IrisInput(BaseModel):
37
+ sepal_length: float = Field(description="Sepal length in cm")
38
+ sepal_width: float = Field(description="Sepal width in cm")
39
+ petal_length: float = Field(description="Petal length in cm")
40
+ petal_width: float = Field(description="Petal width in cm")
41
+
42
+ # Wrap the model
43
+ tool = ModelTool(
44
+ model=clf,
45
+ name="classify_iris",
46
+ description="Classify an iris flower: 0=setosa, 1=versicolor, 2=virginica.",
47
+ input_schema=IrisInput,
48
+ output_name="species",
49
+ output_description="Predicted species index",
50
+ )
51
+
52
+ # Get an OpenAI-ready schema
53
+ import json
54
+ print(json.dumps(tool.to_openai(), indent=2))
55
+
56
+ # Call it directly
57
+ tool.invoke({
58
+ "sepal_length": 5.1, "sepal_width": 3.5,
59
+ "petal_length": 1.4, "petal_width": 0.2,
60
+ })
61
+ # → {"species": 0}
62
+ ```
63
+
64
+ ## Core API
65
+
66
+ ### `ModelTool`
67
+
68
+ ```python
69
+ ModelTool(
70
+ model, # fitted sklearn-compatible estimator
71
+ name: str, # tool name the LLM sees
72
+ description: str, # tool description the LLM sees
73
+ input_schema, # Pydantic BaseModel describing inputs
74
+ output_name: str, # key for the prediction in the returned dict
75
+ output_description: str,
76
+ )
77
+ ```
78
+
79
+ | Method | Returns | What it does |
80
+ |--------|---------|--------------|
81
+ | `.invoke(input_dict)` | `dict` | Validates → predicts → returns `{output_name: value}` |
82
+ | `.to_openai()` | `dict` | OpenAI function-calling schema |
83
+ | `.to_langchain()` | `StructuredTool` | LangChain tool |
84
+ | `.to_callable()` | `Callable` | Plain Python function |
85
+
86
+ ### `ToolRegistry`
87
+
88
+ Group multiple tools for bulk export:
89
+
90
+ ```python
91
+ registry = ToolRegistry([price_tool, risk_tool])
92
+ registry.to_openai() # → list[dict], pass directly to OpenAI
93
+ registry.to_langchain() # → list[StructuredTool]
94
+ registry.get("name") # → ModelTool
95
+ ```
96
+
97
+ ## Field naming rule
98
+
99
+ **Your Pydantic schema field names must exactly match the column names the model was trained on.**
100
+
101
+ predikit maps inputs to features by name, not position. If you trained on a DataFrame with columns `["sqft", "bedrooms"]`, your schema fields must be `sqft` and `bedrooms` — not `sq_ft`, not `Sqft`.
102
+
103
+ ```python
104
+ # ✓ Columns match: sqft, bedrooms, bathrooms
105
+ class GoodInput(BaseModel):
106
+ sqft: float
107
+ bedrooms: float
108
+ bathrooms: float
109
+
110
+ # ✗ Name mismatch — raises ValueError at runtime
111
+ class BadInput(BaseModel):
112
+ square_footage: float # model expects "sqft"
113
+ beds: float # model expects "bedrooms"
114
+ baths: float # model expects "bathrooms"
115
+ ```
116
+
117
+ When there's a mismatch, predikit tells you exactly which names are wrong:
118
+
119
+ ```
120
+ ValueError: Input schema is missing model features: ['sqft', 'bedrooms'].
121
+ Schema has: ['square_footage', 'beds', 'bathrooms'], model expects: ['sqft', 'bedrooms', 'bathrooms']
122
+ ```
123
+
124
+ > **Tip:** If you trained with a numpy array (no DataFrame), predikit has no feature names to check — it uses your schema's field definition order instead.
125
+
126
+ ## Cookbook
127
+
128
+ ### XGBoost regression
129
+
130
+ ```python
131
+ from xgboost import XGBRegressor
132
+ from predikit import ModelTool
133
+
134
+ reg = XGBRegressor().fit(X_train, y_train)
135
+
136
+ class HouseInput(BaseModel):
137
+ sqft: float
138
+ bedrooms: float
139
+ year_built: float
140
+
141
+ tool = ModelTool(
142
+ model=reg,
143
+ name="price_estimate",
144
+ description="Predict home price in USD.",
145
+ input_schema=HouseInput,
146
+ output_name="price_usd",
147
+ output_description="Predicted sale price in USD",
148
+ )
149
+ ```
150
+
151
+ ### Multiple tools in one registry
152
+
153
+ ```python
154
+ registry = ToolRegistry([price_tool, risk_tool, demand_tool])
155
+
156
+ # OpenAI
157
+ response = client.chat.completions.create(
158
+ model="gpt-4o",
159
+ tools=registry.to_openai(),
160
+ ...
161
+ )
162
+
163
+ # LangChain
164
+ agent = initialize_agent(tools=registry.to_langchain(), ...)
165
+ ```
166
+
167
+ ### Bool inputs from an LLM
168
+
169
+ LLMs sometimes return `"yes"`, `"true"`, or `"1"` for boolean fields. predikit coerces these automatically before Pydantic validation:
170
+
171
+ ```python
172
+ class Input(BaseModel):
173
+ has_pool: bool
174
+
175
+ tool.invoke({"has_pool": "yes"}) # → coerced to True
176
+ tool.invoke({"has_pool": "false"}) # → coerced to False
177
+ tool.invoke({"has_pool": "maybe"}) # → raises ValueError with clear message
178
+ ```
179
+
180
+ Supported strings: `true/false`, `yes/no`, `1/0`, `on/off`.
181
+
182
+ ### Orlando real estate demo
183
+
184
+ See [`examples/03_orlando_real_estate.py`](examples/03_orlando_real_estate.py) for a full end-to-end walkthrough: synthetic dataset → XGBoost training → `ModelTool` → registry → OpenAI schema → prediction.
185
+
186
+ ## Roadmap
187
+
188
+ Intentionally out of scope for v0.1 — planned for later releases:
189
+
190
+ - Confidence-aware routing & fallback
191
+ - Multi-model synthesis (agent calls several, reconciles results)
192
+ - MLflow / Snowflake Model Registry integration
193
+ - HuggingFace / PyTorch / TensorFlow support
194
+ - Async invocation
195
+
196
+ ## License
197
+
198
+ MIT © Tejas Tumakuru Ashok
@@ -0,0 +1,45 @@
1
+ """Example 01 — Basic sklearn classifier wrapped as an OpenAI tool."""
2
+ import json
3
+
4
+ from pydantic import BaseModel, Field
5
+ from sklearn.datasets import load_iris
6
+ from sklearn.linear_model import LogisticRegression
7
+
8
+ from predikit import ModelTool
9
+
10
+ # 1. Train
11
+ X, y = load_iris(return_X_y=True)
12
+ clf = LogisticRegression(max_iter=200).fit(X, y)
13
+
14
+
15
+ # 2. Define input schema
16
+ class IrisInput(BaseModel):
17
+ sepal_length: float = Field(description="Sepal length in cm")
18
+ sepal_width: float = Field(description="Sepal width in cm")
19
+ petal_length: float = Field(description="Petal length in cm")
20
+ petal_width: float = Field(description="Petal width in cm")
21
+
22
+
23
+ # 3. Wrap
24
+ tool = ModelTool(
25
+ model=clf,
26
+ name="classify_iris",
27
+ description="Classify an iris flower as setosa (0), versicolor (1), or virginica (2).",
28
+ input_schema=IrisInput,
29
+ output_name="species",
30
+ output_description="Predicted species index: 0=setosa, 1=versicolor, 2=virginica",
31
+ )
32
+
33
+ # 4. OpenAI schema
34
+ print("=== OpenAI Tool Schema ===")
35
+ print(json.dumps(tool.to_openai(), indent=2))
36
+
37
+ # 5. Direct invocation
38
+ print("\n=== Direct Invocation ===")
39
+ result = tool.invoke({"sepal_length": 5.1, "sepal_width": 3.5, "petal_length": 1.4, "petal_width": 0.2})
40
+ print(f"Result: {result}")
41
+
42
+ # 6. Plain callable
43
+ fn = tool.to_callable()
44
+ result2 = fn(sepal_length=6.3, sepal_width=3.3, petal_length=6.0, petal_width=2.5)
45
+ print(f"Via callable: {result2}")
@@ -0,0 +1,54 @@
1
+ """
2
+ Example 02 — XGBoost regression model wrapped as a tool.
3
+ Requires: pip install modelbridge[xgboost]
4
+ """
5
+ import json
6
+
7
+ from pydantic import BaseModel, Field
8
+ from sklearn.datasets import make_regression
9
+ from sklearn.model_selection import train_test_split
10
+
11
+ try:
12
+ from xgboost import XGBRegressor
13
+ except ImportError:
14
+ raise SystemExit("XGBoost not installed. Run: pip install modelbridge[xgboost]")
15
+
16
+ from predikit import ModelTool
17
+
18
+ # 1. Train
19
+ X, y = make_regression(n_samples=500, n_features=4, noise=10, random_state=42)
20
+ X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=42)
21
+ reg = XGBRegressor(n_estimators=100, random_state=42)
22
+ reg.fit(X_train, y_train)
23
+
24
+
25
+ # 2. Define input schema
26
+ class RegressionInput(BaseModel):
27
+ feature_1: float = Field(description="Input feature 1")
28
+ feature_2: float = Field(description="Input feature 2")
29
+ feature_3: float = Field(description="Input feature 3")
30
+ feature_4: float = Field(description="Input feature 4")
31
+
32
+
33
+ # 3. Wrap
34
+ tool = ModelTool(
35
+ model=reg,
36
+ name="xgb_regressor",
37
+ description="Predict a continuous value from 4 numeric features using XGBoost.",
38
+ input_schema=RegressionInput,
39
+ output_name="predicted_value",
40
+ output_description="Predicted numeric output",
41
+ )
42
+
43
+ print("=== OpenAI Tool Schema ===")
44
+ print(json.dumps(tool.to_openai(), indent=2))
45
+
46
+ print("\n=== Invoke ===")
47
+ result = tool.invoke({
48
+ "feature_1": float(X_test[0, 0]),
49
+ "feature_2": float(X_test[0, 1]),
50
+ "feature_3": float(X_test[0, 2]),
51
+ "feature_4": float(X_test[0, 3]),
52
+ })
53
+ print(f"Prediction: {result}")
54
+ print(f"Actual: {y_test[0]:.2f}")
@@ -0,0 +1,122 @@
1
+ """
2
+ Example 03 — Orlando Real Estate Price Predictor (portfolio demo)
3
+
4
+ Trains an XGBoost model on synthetic Orlando-area housing data,
5
+ wraps it as an LLM-callable tool, and shows the full end-to-end flow:
6
+ schema generation → registry export → direct invocation → callable.
7
+
8
+ Requires: pip install modelbridge[xgboost]
9
+ """
10
+ import json
11
+
12
+ import numpy as np
13
+ import pandas as pd
14
+ from pydantic import BaseModel, Field
15
+ from sklearn.model_selection import train_test_split
16
+
17
+ try:
18
+ from xgboost import XGBRegressor
19
+ except ImportError:
20
+ raise SystemExit("XGBoost not installed. Run: pip install modelbridge[xgboost]")
21
+
22
+ from predikit import ModelTool, ToolRegistry
23
+
24
+ # ---------------------------------------------------------------------------
25
+ # 1. Synthetic Orlando dataset
26
+ # ---------------------------------------------------------------------------
27
+ rng = np.random.default_rng(42)
28
+ n = 1_000
29
+
30
+ sqft = rng.integers(800, 4_000, n).astype(float)
31
+ bedrooms = rng.integers(1, 6, n).astype(float)
32
+ bathrooms = rng.choice([1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0], n)
33
+ year_built = rng.integers(1960, 2024, n).astype(float)
34
+ has_pool = rng.integers(0, 2, n).astype(float)
35
+ zip_code_group = rng.integers(0, 5, n).astype(float) # 0=rural … 4=prime
36
+
37
+ price = (
38
+ 80 * sqft
39
+ + 15_000 * bedrooms
40
+ + 10_000 * bathrooms
41
+ + 500 * (year_built - 1960)
42
+ + 20_000 * has_pool
43
+ + 30_000 * zip_code_group
44
+ + rng.normal(0, 15_000, n)
45
+ ).clip(50_000, 1_500_000)
46
+
47
+ df = pd.DataFrame({
48
+ "sqft": sqft,
49
+ "bedrooms": bedrooms,
50
+ "bathrooms": bathrooms,
51
+ "year_built": year_built,
52
+ "has_pool": has_pool,
53
+ "zip_code_group": zip_code_group,
54
+ })
55
+
56
+ X_train, X_test, y_train, y_test = train_test_split(df, price, test_size=0.2, random_state=42)
57
+
58
+ # ---------------------------------------------------------------------------
59
+ # 2. Train model
60
+ # ---------------------------------------------------------------------------
61
+ model = XGBRegressor(n_estimators=200, learning_rate=0.05, random_state=42)
62
+ model.fit(X_train, y_train)
63
+
64
+ # ---------------------------------------------------------------------------
65
+ # 3. Input schema — field names must match DataFrame columns exactly
66
+ # ---------------------------------------------------------------------------
67
+ class OrlandoHouseInput(BaseModel):
68
+ sqft: float = Field(description="Total square footage of the home")
69
+ bedrooms: float = Field(description="Number of bedrooms")
70
+ bathrooms: float = Field(description="Number of bathrooms (0.5 increments)")
71
+ year_built: float = Field(description="Year the home was built")
72
+ has_pool: float = Field(description="1 if the home has a pool, else 0")
73
+ zip_code_group: float = Field(description="Area cluster 0–4 (0=rural, 4=prime location)")
74
+
75
+ # ---------------------------------------------------------------------------
76
+ # 4. Wrap as ModelTool and register
77
+ # ---------------------------------------------------------------------------
78
+ price_tool = ModelTool(
79
+ model=model,
80
+ name="orlando_home_price",
81
+ description=(
82
+ "Estimate the sale price of a residential property in the Orlando, FL metro area "
83
+ "based on its characteristics. Returns predicted price in USD."
84
+ ),
85
+ input_schema=OrlandoHouseInput,
86
+ output_name="estimated_price_usd",
87
+ output_description="Predicted home sale price in US dollars",
88
+ )
89
+
90
+ registry = ToolRegistry([price_tool])
91
+
92
+ # ---------------------------------------------------------------------------
93
+ # 5. Show OpenAI function schema
94
+ # ---------------------------------------------------------------------------
95
+ print("=== OpenAI Function Schema ===")
96
+ print(json.dumps(registry.to_openai(), indent=2))
97
+
98
+ # ---------------------------------------------------------------------------
99
+ # 6. Run a prediction
100
+ # ---------------------------------------------------------------------------
101
+ sample = {
102
+ "sqft": 2_200.0,
103
+ "bedrooms": 3.0,
104
+ "bathrooms": 2.0,
105
+ "year_built": 2005.0,
106
+ "has_pool": 1.0,
107
+ "zip_code_group": 3.0,
108
+ }
109
+
110
+ print("\n=== Sample Prediction ===")
111
+ print(f"Input: {sample}")
112
+ result = price_tool.invoke(sample)
113
+ print(f"Output: {result}")
114
+ print(f"\nEstimated price: ${result['estimated_price_usd']:,.0f}")
115
+
116
+ # ---------------------------------------------------------------------------
117
+ # 7. Same call via plain callable (simulates how an LLM agent would invoke it)
118
+ # ---------------------------------------------------------------------------
119
+ print("\n=== Via Plain Callable (LLM-style) ===")
120
+ fn = price_tool.to_callable()
121
+ result2 = fn(**sample)
122
+ print(f"Result: {result2}")
@@ -0,0 +1,52 @@
1
+ [build-system]
2
+ requires = ["hatchling"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [project]
6
+ name = "predikit"
7
+ version = "0.1.0"
8
+ description = "Turn any trained sklearn/XGBoost model into an LLM-callable tool with auto-generated schemas and typed I/O."
9
+ readme = "README.md"
10
+ license = {text = "MIT"}
11
+ requires-python = ">=3.10"
12
+ authors = [
13
+ {name = "Tejas Tumakuru Ashok", email = "tejas.tumakuruashok@travelandleisure.com"},
14
+ ]
15
+ keywords = ["llm", "agents", "sklearn", "xgboost", "function-calling", "ml-tools"]
16
+ classifiers = [
17
+ "Development Status :: 3 - Alpha",
18
+ "License :: OSI Approved :: MIT License",
19
+ "Programming Language :: Python :: 3",
20
+ "Programming Language :: Python :: 3.10",
21
+ "Programming Language :: Python :: 3.11",
22
+ "Programming Language :: Python :: 3.12",
23
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
24
+ "Intended Audience :: Developers",
25
+ "Intended Audience :: Science/Research",
26
+ ]
27
+ dependencies = [
28
+ "pydantic>=2.0",
29
+ "scikit-learn>=1.2",
30
+ "numpy>=1.24",
31
+ ]
32
+
33
+ [project.optional-dependencies]
34
+ langchain = ["langchain-core>=0.1"]
35
+ xgboost = ["xgboost>=1.7"]
36
+ dev = [
37
+ "pytest>=7.0",
38
+ "pytest-cov>=4.0",
39
+ "xgboost>=1.7",
40
+ "langchain-core>=0.1",
41
+ "pandas>=1.5",
42
+ ]
43
+
44
+ [project.urls]
45
+ Homepage = "https://github.com/Tejas-TA/predikit"
46
+ Repository = "https://github.com/Tejas-TA/predikit"
47
+
48
+ [tool.hatch.build.targets.wheel]
49
+ packages = ["src/predikit"]
50
+
51
+ [tool.pytest.ini_options]
52
+ testpaths = ["tests"]
@@ -0,0 +1,5 @@
1
+ from .tool import ModelTool
2
+ from .registry import ToolRegistry
3
+
4
+ __all__ = ["ModelTool", "ToolRegistry"]
5
+ __version__ = "0.1.0"
@@ -0,0 +1,50 @@
1
+ from typing import Any
2
+ from pydantic import BaseModel
3
+
4
+
5
+ _BOOL_TRUE = {"true", "1", "yes", "on"}
6
+ _BOOL_FALSE = {"false", "0", "no", "off"}
7
+
8
+
9
+ def coerce_value(value: Any, target_type: type) -> Any:
10
+ """Coerce a single value to target_type with LLM-friendly string handling."""
11
+ if isinstance(value, target_type):
12
+ return value
13
+
14
+ if target_type is bool:
15
+ if isinstance(value, str):
16
+ low = value.lower()
17
+ if low in _BOOL_TRUE:
18
+ return True
19
+ if low in _BOOL_FALSE:
20
+ return False
21
+ raise ValueError(f"Cannot interpret {value!r} as bool. Expected one of: true/false, yes/no, 1/0, on/off")
22
+ return bool(value)
23
+
24
+ if target_type is int:
25
+ return int(float(value))
26
+
27
+ if target_type is float:
28
+ return float(value)
29
+
30
+ if target_type is str:
31
+ return str(value)
32
+
33
+ return value
34
+
35
+
36
+ def coerce_inputs(validated: BaseModel, meta: dict) -> list:
37
+ """Return feature values in the order the model expects them."""
38
+ data = validated.model_dump()
39
+ feature_names = meta.get("feature_names")
40
+ if feature_names:
41
+ missing = [f for f in feature_names if f not in data]
42
+ if missing:
43
+ raise ValueError(
44
+ f"Input schema is missing model features: {missing}. "
45
+ f"Field names in your Pydantic schema must exactly match the column names "
46
+ f"used during model training. "
47
+ f"Schema has: {list(data.keys())}, model expects: {feature_names}"
48
+ )
49
+ return [data[f] for f in feature_names]
50
+ return list(data.values())
File without changes
@@ -0,0 +1,25 @@
1
+ from __future__ import annotations
2
+ from typing import TYPE_CHECKING
3
+
4
+ if TYPE_CHECKING:
5
+ from ..tool import ModelTool
6
+
7
+
8
+ def to_langchain_tool(tool: ModelTool):
9
+ """Convert a ModelTool to a LangChain StructuredTool."""
10
+ try:
11
+ from langchain_core.tools import StructuredTool
12
+ except ImportError:
13
+ raise ImportError(
14
+ "langchain-core is required. Install with: pip install predikit[langchain]"
15
+ )
16
+
17
+ def _run(**kwargs) -> dict:
18
+ return tool.invoke(kwargs)
19
+
20
+ return StructuredTool.from_function(
21
+ func=_run,
22
+ name=tool.name,
23
+ description=tool.description,
24
+ args_schema=tool.input_schema,
25
+ )
@@ -0,0 +1,20 @@
1
+ from __future__ import annotations
2
+ from typing import TYPE_CHECKING
3
+
4
+ if TYPE_CHECKING:
5
+ from ..tool import ModelTool
6
+
7
+
8
+ def to_openai_schema(tool: ModelTool) -> dict:
9
+ """Convert a ModelTool to an OpenAI function-calling schema dict."""
10
+ schema = tool.input_schema.model_json_schema()
11
+ schema.pop("title", None)
12
+
13
+ return {
14
+ "type": "function",
15
+ "function": {
16
+ "name": tool.name,
17
+ "description": tool.description,
18
+ "parameters": schema,
19
+ },
20
+ }
@@ -0,0 +1,22 @@
1
+ from typing import Any
2
+
3
+
4
+ def introspect(model: Any) -> dict:
5
+ """Extract metadata from a fitted sklearn-compatible estimator."""
6
+ meta: dict = {}
7
+
8
+ meta["feature_names"] = (
9
+ list(model.feature_names_in_) if hasattr(model, "feature_names_in_") else None
10
+ )
11
+ meta["n_features"] = (
12
+ int(model.n_features_in_) if hasattr(model, "n_features_in_") else None
13
+ )
14
+
15
+ if hasattr(model, "classes_"):
16
+ meta["task"] = "classification"
17
+ meta["classes"] = list(model.classes_)
18
+ else:
19
+ meta["task"] = "regression"
20
+ meta["classes"] = None
21
+
22
+ return meta
@@ -0,0 +1,20 @@
1
+ from __future__ import annotations
2
+ from .tool import ModelTool
3
+
4
+
5
+ class ToolRegistry:
6
+ """Bundles multiple ModelTools and provides bulk export methods."""
7
+
8
+ def __init__(self, tools: list[ModelTool]) -> None:
9
+ self._tools: dict[str, ModelTool] = {t.name: t for t in tools}
10
+
11
+ def get(self, name: str) -> ModelTool:
12
+ if name not in self._tools:
13
+ raise KeyError(f"No tool named '{name}'. Available: {list(self._tools)}")
14
+ return self._tools[name]
15
+
16
+ def to_openai(self) -> list[dict]:
17
+ return [t.to_openai() for t in self._tools.values()]
18
+
19
+ def to_langchain(self) -> list:
20
+ return [t.to_langchain() for t in self._tools.values()]
@@ -0,0 +1,93 @@
1
+ from __future__ import annotations
2
+ from typing import Any, Callable
3
+
4
+ import numpy as np
5
+ from pydantic import BaseModel
6
+
7
+ from .introspect import introspect
8
+ from .coerce import coerce_inputs, coerce_value
9
+ from .exporters.openai import to_openai_schema
10
+ from .exporters.langchain import to_langchain_tool
11
+
12
+ # Only apply pre-coercion for these scalar types; everything else goes straight to Pydantic.
13
+ _SCALAR_TYPES = (bool, int, float, str)
14
+
15
+
16
+ class ModelTool:
17
+ """Wraps a fitted sklearn-compatible model as an LLM-callable tool."""
18
+
19
+ def __init__(
20
+ self,
21
+ model: Any,
22
+ name: str,
23
+ description: str,
24
+ input_schema: type[BaseModel],
25
+ output_name: str,
26
+ output_description: str,
27
+ ) -> None:
28
+ self.model = model
29
+ self.name = name
30
+ self.description = description
31
+ self.input_schema = input_schema
32
+ self.output_name = output_name
33
+ self.output_description = output_description
34
+ self._meta = introspect(model)
35
+
36
+ def invoke(self, input_dict: dict) -> dict:
37
+ """Validate inputs, run prediction, return {output_name: value}."""
38
+ pre_coerced = self._pre_coerce(input_dict)
39
+ try:
40
+ validated = self.input_schema(**pre_coerced)
41
+ except Exception as exc:
42
+ raise ValueError(f"Input validation failed for '{self.name}': {exc}") from exc
43
+
44
+ features = coerce_inputs(validated, self._meta)
45
+ X = self._to_array(features)
46
+ prediction = self.model.predict(X)[0]
47
+
48
+ if hasattr(prediction, "item"):
49
+ prediction = prediction.item()
50
+
51
+ return {self.output_name: prediction}
52
+
53
+ def to_openai(self) -> dict:
54
+ """Return an OpenAI function-calling schema dict."""
55
+ return to_openai_schema(self)
56
+
57
+ def to_langchain(self):
58
+ """Return a LangChain StructuredTool."""
59
+ return to_langchain_tool(self)
60
+
61
+ def to_callable(self) -> Callable[..., dict]:
62
+ """Return a plain Python function that calls invoke()."""
63
+ def _fn(**kwargs) -> dict:
64
+ return self.invoke(kwargs)
65
+ _fn.__name__ = self.name
66
+ _fn.__doc__ = self.description
67
+ return _fn
68
+
69
+ def _pre_coerce(self, input_dict: dict) -> dict:
70
+ # Runs coerce_value before Pydantic so LLM strings like "yes"/"no" work for bool fields.
71
+ fields = self.input_schema.model_fields
72
+ result = {}
73
+ for k, v in input_dict.items():
74
+ field = fields.get(k)
75
+ annotation = getattr(field, "annotation", None) if field else None
76
+ if annotation in _SCALAR_TYPES:
77
+ try:
78
+ result[k] = coerce_value(v, annotation)
79
+ except (ValueError, TypeError):
80
+ result[k] = v # let Pydantic surface the error with full context
81
+ else:
82
+ result[k] = v
83
+ return result
84
+
85
+ def _to_array(self, features: list) -> Any:
86
+ feature_names = self._meta.get("feature_names")
87
+ if feature_names:
88
+ try:
89
+ import pandas as pd
90
+ return pd.DataFrame([dict(zip(feature_names, features))])
91
+ except ImportError:
92
+ pass
93
+ return np.array([features])
File without changes
@@ -0,0 +1,57 @@
1
+ import pytest
2
+ from pydantic import BaseModel
3
+
4
+ from predikit.coerce import coerce_inputs, coerce_value
5
+
6
+
7
+ def test_coerce_bool_strings():
8
+ assert coerce_value("true", bool) is True
9
+ assert coerce_value("yes", bool) is True
10
+ assert coerce_value("false", bool) is False
11
+ assert coerce_value("0", bool) is False
12
+
13
+
14
+ def test_coerce_bool_invalid():
15
+ with pytest.raises(ValueError):
16
+ coerce_value("maybe", bool)
17
+
18
+
19
+ def test_coerce_int_from_string():
20
+ assert coerce_value("42", int) == 42
21
+ assert coerce_value("3.7", int) == 3
22
+
23
+
24
+ def test_coerce_float_from_string():
25
+ assert coerce_value("3.14", float) == pytest.approx(3.14)
26
+
27
+
28
+ def test_coerce_inputs_ordered_by_feature_names():
29
+ class Schema(BaseModel):
30
+ b: float
31
+ a: float
32
+
33
+ validated = Schema(b=2.0, a=1.0)
34
+ meta = {"feature_names": ["a", "b"]}
35
+ result = coerce_inputs(validated, meta)
36
+ assert result == [1.0, 2.0]
37
+
38
+
39
+ def test_coerce_inputs_no_feature_names():
40
+ class Schema(BaseModel):
41
+ x: float
42
+ y: float
43
+
44
+ validated = Schema(x=3.0, y=4.0)
45
+ meta = {"feature_names": None}
46
+ result = coerce_inputs(validated, meta)
47
+ assert result == [3.0, 4.0]
48
+
49
+
50
+ def test_coerce_inputs_missing_feature_raises():
51
+ class Schema(BaseModel):
52
+ x: float
53
+
54
+ validated = Schema(x=1.0)
55
+ meta = {"feature_names": ["x", "y"]}
56
+ with pytest.raises(ValueError, match="missing"):
57
+ coerce_inputs(validated, meta)
@@ -0,0 +1,42 @@
1
+ from pydantic import BaseModel, Field
2
+ from sklearn.datasets import load_iris
3
+ from sklearn.linear_model import LogisticRegression
4
+
5
+ from predikit import ModelTool
6
+
7
+
8
+ class IrisInput(BaseModel):
9
+ sepal_length: float = Field(description="Sepal length in cm")
10
+ sepal_width: float = Field(description="Sepal width in cm")
11
+ petal_length: float = Field(description="Petal length in cm")
12
+ petal_width: float = Field(description="Petal width in cm")
13
+
14
+
15
+ def _make_tool() -> ModelTool:
16
+ X, y = load_iris(return_X_y=True)
17
+ clf = LogisticRegression(max_iter=200).fit(X, y)
18
+ return ModelTool(
19
+ model=clf,
20
+ name="iris_classifier",
21
+ description="Classify iris species from petal/sepal measurements",
22
+ input_schema=IrisInput,
23
+ output_name="species",
24
+ output_description="Predicted species index",
25
+ )
26
+
27
+
28
+ def test_openai_schema_top_level_structure():
29
+ schema = _make_tool().to_openai()
30
+ assert schema["type"] == "function"
31
+ assert schema["function"]["name"] == "iris_classifier"
32
+ assert "parameters" in schema["function"]
33
+
34
+
35
+ def test_openai_schema_has_all_fields():
36
+ props = _make_tool().to_openai()["function"]["parameters"]["properties"]
37
+ assert set(props.keys()) == {"sepal_length", "sepal_width", "petal_length", "petal_width"}
38
+
39
+
40
+ def test_openai_schema_no_title_at_top():
41
+ params = _make_tool().to_openai()["function"]["parameters"]
42
+ assert "title" not in params
@@ -0,0 +1,31 @@
1
+ from sklearn.datasets import load_iris
2
+ from sklearn.linear_model import LinearRegression, LogisticRegression
3
+
4
+ from predikit.introspect import introspect
5
+
6
+
7
+ def test_classifier_metadata():
8
+ X, y = load_iris(return_X_y=True, as_frame=True)
9
+ clf = LogisticRegression(max_iter=200).fit(X, y)
10
+ meta = introspect(clf)
11
+ assert meta["task"] == "classification"
12
+ assert meta["classes"] == [0, 1, 2]
13
+ assert meta["feature_names"] == list(X.columns)
14
+ assert meta["n_features"] == 4
15
+
16
+
17
+ def test_regressor_metadata():
18
+ X, y = load_iris(return_X_y=True, as_frame=True)
19
+ reg = LinearRegression().fit(X, y)
20
+ meta = introspect(reg)
21
+ assert meta["task"] == "regression"
22
+ assert meta["classes"] is None
23
+ assert meta["feature_names"] == list(X.columns)
24
+
25
+
26
+ def test_no_feature_names():
27
+ X, y = load_iris(return_X_y=True)
28
+ reg = LinearRegression().fit(X, y)
29
+ meta = introspect(reg)
30
+ assert meta["feature_names"] is None
31
+ assert meta["n_features"] == 4
@@ -0,0 +1,45 @@
1
+ import pytest
2
+ from pydantic import BaseModel
3
+ from sklearn.datasets import load_iris
4
+ from sklearn.linear_model import LogisticRegression
5
+
6
+ from predikit import ModelTool, ToolRegistry
7
+
8
+
9
+ class IrisInput(BaseModel):
10
+ sepal_length: float
11
+ sepal_width: float
12
+ petal_length: float
13
+ petal_width: float
14
+
15
+
16
+ @pytest.fixture
17
+ def registry():
18
+ X, y = load_iris(return_X_y=True)
19
+ clf = LogisticRegression(max_iter=200).fit(X, y)
20
+ tool = ModelTool(
21
+ model=clf,
22
+ name="iris_classifier",
23
+ description="Classify iris species",
24
+ input_schema=IrisInput,
25
+ output_name="species",
26
+ output_description="Predicted species",
27
+ )
28
+ return ToolRegistry([tool])
29
+
30
+
31
+ def test_get_tool(registry):
32
+ tool = registry.get("iris_classifier")
33
+ assert tool.name == "iris_classifier"
34
+
35
+
36
+ def test_get_missing_raises(registry):
37
+ with pytest.raises(KeyError):
38
+ registry.get("nonexistent")
39
+
40
+
41
+ def test_to_openai_returns_list(registry):
42
+ schemas = registry.to_openai()
43
+ assert isinstance(schemas, list)
44
+ assert len(schemas) == 1
45
+ assert schemas[0]["type"] == "function"
@@ -0,0 +1,84 @@
1
+ import numpy as np
2
+ import pytest
3
+ from pydantic import BaseModel
4
+ from sklearn.datasets import load_iris
5
+ from sklearn.linear_model import LogisticRegression
6
+
7
+ from predikit import ModelTool
8
+
9
+
10
+ class IrisInput(BaseModel):
11
+ sepal_length: float
12
+ sepal_width: float
13
+ petal_length: float
14
+ petal_width: float
15
+
16
+
17
+ @pytest.fixture
18
+ def iris_tool():
19
+ X, y = load_iris(return_X_y=True)
20
+ clf = LogisticRegression(max_iter=200).fit(X, y)
21
+ return ModelTool(
22
+ model=clf,
23
+ name="iris_classifier",
24
+ description="Classify iris species from measurements",
25
+ input_schema=IrisInput,
26
+ output_name="species",
27
+ output_description="Predicted iris species (0=setosa, 1=versicolor, 2=virginica)",
28
+ )
29
+
30
+
31
+ def test_invoke_returns_dict(iris_tool):
32
+ result = iris_tool.invoke(
33
+ {"sepal_length": 5.1, "sepal_width": 3.5, "petal_length": 1.4, "petal_width": 0.2}
34
+ )
35
+ assert "species" in result
36
+ assert result["species"] in [0, 1, 2]
37
+
38
+
39
+ def test_invoke_string_inputs_coerced(iris_tool):
40
+ result = iris_tool.invoke(
41
+ {"sepal_length": "5.1", "sepal_width": "3.5", "petal_length": "1.4", "petal_width": "0.2"}
42
+ )
43
+ assert result["species"] in [0, 1, 2]
44
+
45
+
46
+ def test_invoke_invalid_input_raises(iris_tool):
47
+ with pytest.raises(ValueError):
48
+ iris_tool.invoke({"sepal_length": "not_a_number"})
49
+
50
+
51
+ def test_to_callable(iris_tool):
52
+ fn = iris_tool.to_callable()
53
+ assert callable(fn)
54
+ result = fn(sepal_length=5.1, sepal_width=3.5, petal_length=1.4, petal_width=0.2)
55
+ assert "species" in result
56
+
57
+
58
+ def test_output_is_native_python_type(iris_tool):
59
+ result = iris_tool.invoke(
60
+ {"sepal_length": 5.1, "sepal_width": 3.5, "petal_length": 1.4, "petal_width": 0.2}
61
+ )
62
+ assert isinstance(result["species"], int)
63
+
64
+
65
+ def test_invoke_bool_string_coerced_through_invoke():
66
+ class FlagInput(BaseModel):
67
+ value: float
68
+ active: bool
69
+
70
+ X = np.array([[1.0, 0], [2.0, 1], [3.0, 0], [4.0, 1]])
71
+ y = np.array([0, 1, 0, 1])
72
+ clf = LogisticRegression().fit(X, y)
73
+
74
+ tool = ModelTool(
75
+ model=clf,
76
+ name="flag_test",
77
+ description="test bool coercion",
78
+ input_schema=FlagInput,
79
+ output_name="result",
80
+ output_description="predicted class",
81
+ )
82
+ assert tool.invoke({"value": 3.0, "active": "yes"})["result"] in [0, 1]
83
+ assert tool.invoke({"value": 1.0, "active": "false"})["result"] in [0, 1]
84
+ assert tool.invoke({"value": 3.0, "active": "on"})["result"] in [0, 1]