predikit

Tejas-TAMIT活跃维护

把训练好的 sklearn、XGBoost 等 ML 模型直接暴露成 LLM Agent 可调用的工具,省去自己写胶水代码,让 Agent 能调用传统模型做预测。

The missing bridge between your ML models and your AI agents.

v0.6.4发布于 2026-10-28 18:34 UTC382star128fork同步于 2026-11-10 12:00 UTC

怎么上手

由 README 提炼,执行前请确认
pip install predikit
硬件要求按 README 说明

README 未说明,仅要求 Python 3.10+,无 GPU 或显存要求

下载:源码 zip·GitHub Releases(安装包)本站不镜像安装包

项目说明

由 README 自动翻译提炼

把训练好的 sklearn/XGBoost 模型包装成 LLM 可调用的工具,自动生成 JSON Schema 与类型校验

把训练好的 sklearn、XGBoost 等 ML 模型直接暴露成 LLM Agent 可调用的工具,省去自己写胶水代码,让 Agent 能调用传统模型做预测。

核心能力

  • 用 Pydantic BaseModel 自动生成工具 JSON Schema
  • 一行 to_openai / to_langchain 导出给 LLM 框架
  • invoke 前做 Pydantic v2 类型校验并给出清晰报错
  • ModelEnsemble 提供 collect/mean/vote 等 5 种聚合策略
  • 支持 confidence_threshold 低置信度告警、报错或转 fallback

容易踩的坑

XGBoost、LangChain、MLflow、Snowflake、CLI、MCP 均为可选依赖,需按需安装对应 extra;Pydantic 字段名必须与训练时 DataFrame 列名完全一致,否则运行时报 ValueError;confidence_threshold 仅对实现 predict_proba 的分类器生效,回归模型不受影响

内容由模型根据 README 生成,可能不完整或不准确,以官方文档为准。看原文 README

最近发版

5 个版本v0.6.4
  1. v0.6.4最新正式版2026-10-28 · 2026-10-28 18:34 UTC

    Validate unnamed model feature counts when creating a ModelTool, producing an actionable schema/model mismatch error before the first prediction instead of exposing an estimator error.

  2. v0.6.32026-09-17 · 2026-09-17 15:26 UTC
    查看更新说明

    Minor bug fixes and performance improvements

  3. v0.6.22026-08-22 · 2026-08-22 19:10 UTC
    查看更新说明

    Highlights

    py.typed is now shipped. predikit has always been fully annotated and mypy-clean, and has always declared the Typing :: Typed classifier — but the PEP 561 marker file was missing, so type checkers in downstream projects silently ignored every annotation in the package. If you use mypy or Pyright against predikit, you get real type information for the first time in this release. No API changed; this is purely additive.

    Added

    • py.typed marker file, so type checkers consume the package's annotations (PEP 561). Verified present in the built wheel.

    Changed

    • from_snowflake() now imports snowflake.ml.registry inside the function body rather than at module level. This matches from_mlflow() and the project's own rule that optional dependencies are imported lazily. The module-level Registry = None sentinel and its type: ignore[assignment,misc] are gone, and the ImportError now chains the underlying cause via raise ... from exc.
    • predikit inspect builds its synthetic input schema with dict.fromkeys instead of a dict comprehension.
    • predikit/__init__.py gained a module docstring, and __all__ is now sorted.
    • CI lints and format-checks examples/ alongside src/ and tests/. The seven example scripts are user-facing documentation and were previously outside the gate; they were already clean, so this locks in the current state against future drift.

    Removed

    • uv.lock (5,586 lines). Nothing consumed it — CI installs with pip install -e ".[dev]", pre-commit doesn't use uv, and the contributor docs never referenced it. A library also shouldn't pin its consumers' transitive dependencies. This drops roughly a third of the repository by size.
    • .claude/settings.local.json from version control. It held per-developer machine paths; it's now gitignored.

    Internal

    • tests/test_loaders_snowflake.py: the six tests that patched predikit.loaders.snowflake.Registry now go through a mock_registry fixture that injects a stand-in snowflake.ml.registry into sys.modules. This exercises the real import path instead of a module attribute that only existed because of the module-level import, and the ImportError test no longer needs importlib.reload.
    • tests/test_logging.py rewritten to match the style of the other test modules, plus a new test_quiet_by_default case — nothing previously asserted that verbose=False stays silent.
    • .gitignore reorganized with section headers. Dropped a stale file.py entry; added .mypy_cache/, .ruff_cache/, .claude/settings.local.json, and *.joblib / *.pkl (written by the examples and CLI tests).
    • CONTRIBUTING.md and CLAUDE.md updated for the widened lint scope.

    Verification

    • 110 tests passing (8 integration tests deselected)
    • ruff check and ruff format --check clean across src/, tests/, examples/
    • mypy src/predikit — no issues in 15 source files
    • Wheel builds with predikit/py.typed included

    Install

    pip install --upgrade predikit
    

    Full Changelog: https://github.com/Tejas-TA/predikit/compare/v0.6.1...v0.6.2

  4. v0.6.12026-08-15 · 2026-08-15 15:31 UTC
    查看更新说明

    Highlights MCP tools shipped in 0.6.0 advertised no output schema. This release fixes that, adds host/port control for the HTTP transport, and puts the MCP integration under test against the real SDK.

    pip install --upgrade predikit

    🐛 Fixed 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. 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. 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 surfacing as Could not load registry target. coerce_value() converts True / False to 1 / 0 for int fields. isinstance(True, int) is true in Python, so bools previously passed through unconverted.

    ✨ Added --host and --port 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.

    predikit serve myapp.tools:registry --transport streamable-http --host 0.0.0.0 --port 9000 ToolRegistry.items(), .names(), iter, and len as the public way to enumerate a registry. create_mcp_server() previously used the private _tools / _ensembles attributes; to_openai() and to_langchain() now share the same accessor.

    Python 3.13 in the CI matrix and the package classifiers.

    🔧 Changed 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 100% coverage of mcp.py while never exercising the SDK — which is how the missing output schema shipped. Test coverage for predikit serve, which previously had none.

  5. v0.6.02026-08-15 · 2026-08-15 14:45 UTC
    查看更新说明

用户点评

登录用户的真实使用记录,可按版本和硬件筛选
共 0 条

    还没有人点评。第一个说出真实体验的人最有价值。

    README 原文

    未翻译
    tool = ModelTool(model=clf, name="classify_iris", ...)
    tool.to_openai()              # OpenAI function schema, ready to pass to the API
    tool.invoke({"sqft": 2200})   # → {"price_usd": 370730}
    

    Table of Contents


    Why predikit?

    Most ML pipelines stop at .predict(). Getting that prediction callable by an LLM agent — with validated inputs, correct types, and a schema the model can reason about — requires glue code that's tedious to write and easy to get wrong.

    predikit handles that layer for you:

    Without predikitWith predikit
    SchemaHand-write JSON Schema for every modelAuto-generated from your Pydantic BaseModel
    Type safetyManual casting, silent failuresPydantic v2 validation with clear error messages
    LLM integrationOpenAI / LangChain boilerplate per model.to_openai() / .to_langchain() in one line
    Ensemble routingCustom aggregation logic per projectModelEnsemble with 5 built-in strategies
    Confidence handlingWrite your own threshold checksconfidence_threshold + on_low_confidence
    Model registriesManual MLflow / Snowflake registry callsfrom_mlflow() / from_snowflake() loaders
    Asyncasyncio.get_event_loop().run_in_executor(...)await tool.ainvoke(inputs)

    Works with

    predikit is designed to plug into the tools your team already uses:

    Models — scikit-learn · XGBoost

    LLM frameworks — OpenAI function calling · LangChain · any tool-calling API that accepts JSON Schema

    Model registries — MLflow Model Registry · Snowflake Model Registry

    Validation — Pydantic v2

    Async — asyncio-compatible via ainvoke()


    Install

    pip install predikit
    
    # Optional extras
    pip install predikit[xgboost]    # XGBoost support
    pip install predikit[langchain]  # LangChain StructuredTool export
    pip install predikit[mlflow]     # MLflow Model Registry loader
    pip install predikit[snowflake]  # Snowflake Model Registry loader
    pip install predikit[cli]        # `predikit inspect` and `predikit serve`
    pip install predikit[mcp,cli]    # serve a registry over MCP
    

    Quick start

    from pydantic import BaseModel, Field
    from sklearn.datasets import load_iris
    from sklearn.linear_model import LogisticRegression
    from predikit import ModelTool
    
    # Train
    X, y = load_iris(return_X_y=True)
    clf = LogisticRegression(max_iter=200).fit(X, y)
    
    # Define what the LLM will pass in
    class IrisInput(BaseModel):
        sepal_length: float = Field(description="Sepal length 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")
    
    # Wrap the model
    tool = ModelTool(
        model=clf,
        name="classify_iris",
        description="Classify an iris flower: 0=setosa, 1=versicolor, 2=virginica.",
        input_schema=IrisInput,
        output_name="species",
        output_description="Predicted species index",
    )
    
    # Get an OpenAI-ready schema
    import json
    print(json.dumps(tool.to_openai(), indent=2))
    
    # Call it directly
    tool.invoke({
        "sepal_length": 5.1, "sepal_width": 3.5,
        "petal_length": 1.4, "petal_width": 0.2,
    })
    # → {"species": 0}
    

    Core API

    ModelTool

    ModelTool(
        model,               # fitted sklearn-compatible estimator
        name: str,           # tool name the LLM sees
        description: str,    # tool description the LLM sees
        input_schema,        # Pydantic BaseModel describing inputs
        output_name: str,    # key for the prediction in the returned dict
        output_description: str,
    )
    
    MethodReturnsWhat it does
    .invoke(input_dict)dictValidates → predicts → returns {output_name: value}
    .ainvoke(input_dict)dictAsync version of .invoke()
    .to_openai()dictOpenAI function-calling schema
    .to_langchain()StructuredToolLangChain tool
    .to_callable()CallablePlain Python function

    ToolRegistry

    Group multiple tools for bulk export:

    registry = ToolRegistry([price_tool, risk_tool])
    registry.to_openai()     # → list[dict], pass directly to OpenAI
    registry.to_langchain()  # → list[StructuredTool]
    registry.get("name")     # → ModelTool
    

    ModelEnsemble

    Call multiple models and reconcile their outputs in one step:

    ModelEnsemble(
        tools: list[ModelTool],   # models to run in parallel
        name: str,                # ensemble tool name the LLM sees
        description: str,
        strategy: str,            # "collect" | "mean" | "vote" | "weighted_mean" | "weighted_vote"
        weights: list[float],     # optional, for weighted strategies
    )
    
    StrategyBehaviour
    "collect"Merges all outputs into one dict (tools can have different output_name)
    "mean"Averages numeric outputs (all tools must share output_name)
    "vote"Majority class vote (all tools must share output_name)
    "weighted_mean"Weighted average — provide a weights list
    "weighted_vote"Weighted majority vote — provide a weights list

    ModelEnsemble exposes the same .invoke(), .ainvoke(), .to_openai(), and .to_langchain() interface as ModelTool.


    Field naming rule

    Your Pydantic schema field names must exactly match the column names the model was trained on.

    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.

    # ✓ Columns match: sqft, bedrooms, bathrooms
    class GoodInput(BaseModel):
        sqft:      float
        bedrooms:  float
        bathrooms: float
    
    # ✗ Name mismatch — raises ValueError at runtime
    class BadInput(BaseModel):
        square_footage: float  # model expects "sqft"
        beds:           float  # model expects "bedrooms"
        baths:          float  # model expects "bathrooms"
    

    When there's a mismatch, predikit tells you exactly which names are wrong:

    ValueError: Input schema is missing model features: ['sqft', 'bedrooms'].
    Schema has: ['square_footage', 'beds', 'bathrooms'], model expects: ['sqft', 'bedrooms', 'bathrooms']
    

    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.


    Cookbook

    XGBoost regression

    from xgboost import XGBRegressor
    from predikit import ModelTool
    
    reg = XGBRegressor().fit(X_train, y_train)
    
    class HouseInput(BaseModel):
        sqft:       float
        bedrooms:   float
        year_built: float
    
    tool = ModelTool(
        model=reg,
        name="price_estimate",
        description="Predict home price in USD.",
        input_schema=HouseInput,
        output_name="price_usd",
        output_description="Predicted sale price in USD",
    )
    

    Multiple tools in one registry

    registry = ToolRegistry([price_tool, risk_tool, demand_tool])
    
    # OpenAI
    response = client.chat.completions.create(
        model="gpt-4o",
        tools=registry.to_openai(),
        ...
    )
    
    # LangChain
    agent = initialize_agent(tools=registry.to_langchain(), ...)
    

    MCP server

    Expose the same registry to MCP-compatible clients with the optional MCP extra:

    pip install predikit[mcp,cli]
    predikit serve my_tools:registry
    

    my_tools:registry must point to a ToolRegistry, or to a zero-argument function that returns one. The default transport is stdio; use --transport streamable-http when serving over HTTP, optionally with --host and --port. MCP tools advertise their output schemas and use the same Pydantic input validation and model execution as direct invoke() calls. For an in-process server instead of the CLI, install predikit[mcp] and call create_mcp_server(registry).

    Bool inputs from an LLM

    LLMs sometimes return "yes", "true", or "1" for boolean fields. predikit coerces these automatically before Pydantic validation:

    class Input(BaseModel):
        has_pool: bool
    
    tool.invoke({"has_pool": "yes"})   # → coerced to True
    tool.invoke({"has_pool": "false"}) # → coerced to False
    tool.invoke({"has_pool": "maybe"}) # → raises ValueError with clear message
    

    Supported strings: true/false, yes/no, 1/0, on/off.

    Confidence-aware routing

    Route uncertain predictions to a fallback tool, or raise an error the agent can catch:

    from predikit import ModelTool, LowConfidenceError
    
    tool = ModelTool(
        model=clf,
        name="churn_risk",
        description="Predict member churn risk.",
        input_schema=MemberInput,
        output_name="churn_probability",
        output_description="Probability of churn (0–1)",
        confidence_threshold=0.80,       # classifiers with predict_proba only
        on_low_confidence="warn",        # "warn" | "raise" | "fallback"
        fallback_tool=rule_based_tool,   # used when mode="fallback"
    )
    
    result = tool.invoke(inputs)
    if result.get("_low_confidence"):
        print(f"Uncertain ({result['_confidence']:.2f}) — consider routing to a human")
    
    modebehaviour
    "warn"returns prediction + _confidence + _low_confidence: True
    "raise"raises LowConfidenceError
    "fallback"invokes fallback_tool and returns its result

    Only applies to classifiers that implement predict_proba. Regressors are unaffected.

    Multi-model ensemble

    from predikit import ModelEnsemble, ToolRegistry
    
    ensemble = ModelEnsemble(
        tools=[price_tool_a, price_tool_b],
        name="averaged_price",
        description="Ensemble price: mean of two XGBoost models.",
        strategy="mean",              # "collect" | "mean" | "vote"
    )
    
    result  = ensemble.invoke(inputs)  # → {"price_usd": 370112}
    schema  = ensemble.to_openai()     # works exactly like ModelTool
    

    Register ensembles alongside individual tools:

    registry = ToolRegistry(tools=[price_tool], ensembles=[ensemble])
    registry.to_openai()  # includes both tools and ensembles
    

    MLflow Model Registry loader

    Load a registered MLflow model directly — no manual .load_model() call:

    from predikit.loaders import from_mlflow
    
    tool = from_mlflow(
        model_uri="models:/churn-classifier/Production",
        name="churn_risk",
        description="Predict member churn probability.",
        input_schema=MemberInput,
        output_name="churn_probability",
        output_description="Churn probability 0–1",
    )
    
    tool.invoke({"tenure_months": 24, "trips_last_year": 2, "avg_spend": 500})
    # → {"churn_probability": 0.73}
    

    The loader auto-detects classes_ and feature_names_in_ from the underlying sklearn model, so confidence routing and ensemble work unchanged. Requires pip install predikit[mlflow].

    Snowflake Model Registry loader

    Load a model registered in the Snowflake Model Registry via the Snowpark ML Python library:

    from predikit.loaders import from_snowflake
    
    tool = from_snowflake(
        session=snowpark_session,
        model_name="VACATION_CHURN",
        model_version="V3",
        name="churn_risk",
        description="Churn classifier.",
        input_schema=MemberInput,
        output_name="churn_probability",
        output_description="Churn probability 0–1",
        output_method="predict",   # method to call on the Snowflake model object
    )
    

    Pass output_method="predict_proba" or any other method your Snowflake model exposes. The returned ModelTool is identical to one built directly — all exporters, confidence routing, and ensemble strategies work as-is. Requires pip install predikit[snowflake].

    End-to-end demo

    See examples/03_orlando_real_estate.py for a full walkthrough: synthetic dataset → XGBoost training → ModelTool → registry → OpenAI schema → prediction.


    Roadmap

    Shipped

    • OpenAI function-calling schema export
    • LangChain StructuredTool export
    • Pydantic v2 typed I/O validation
    • Multi-model ModelEnsemble — collect / mean / vote / weighted variants
    • Confidence routing — warn / raise / fallback
    • Async ainvoke() via thread pool executor
    • MLflow Model Registry loader
    • Snowflake Model Registry loader
    • predikit inspect CLI

    Planned

    • HuggingFace / PyTorch / TensorFlow model support
    • Streaming inference support
    • OpenAI Assistants API integration
    • MCP server mode (stdio and Streamable HTTP)

    Contributing

    See CONTRIBUTING.md for development setup, code style, and PR guidelines. The CHANGELOG tracks notable changes per release.

    Issues and PRs are welcome — if you're wrapping a model type or registry that predikit doesn't support yet, open a discussion.

    License

    MIT © Tejas Tumakuru Ashok