Source code for process_improve.tool_spec
"""(c) Kevin Dunn, 2010-2026. MIT License.
Tool-call-first infrastructure for process-improve.
Provides the ``@tool_spec`` decorator, a global registry of all decorated
functions, and helpers for Anthropic-compatible tool-use integrations.
Quick start
-----------
Import the decorated tools and pass the specs to the Anthropic client::
import anthropic
from process_improve.tool_spec import get_tool_specs, execute_tool_call
client = anthropic.Anthropic()
response = client.messages.create(
model="claude-opus-4-6",
max_tokens=1024,
tools=get_tool_specs(),
messages=[{"role": "user", "content": "Are there outliers in [1,2,3,100]?"}],
)
# Dispatch tool calls from the response
for block in response.content:
if block.type == "tool_use":
result = execute_tool_call(block.name, block.input)
"""
from __future__ import annotations
import logging
import math
from collections.abc import Callable
from typing import Any
import numpy as np
from pydantic import BaseModel, ValidationError
logger = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
# Registry
# ---------------------------------------------------------------------------
#: Maps tool name -> decorated callable. Populated by ``@tool_spec``.
_TOOL_REGISTRY: dict[str, Callable[..., Any]] = {}
#: Whether ``discover_tools()`` has already run.
_discovery_done: bool = False
# ---------------------------------------------------------------------------
# RNG metadata validation
# ---------------------------------------------------------------------------
def _validate_rng_metadata(name: str, rng: dict[str, Any]) -> None:
"""Reject ``rng`` payloads that don't match the published contract.
The contract is intentionally narrow so downstream consumers (the
factorial reproducible-export service) can introspect specs without
defensive type-checking. See ``CLAUDE.md`` for the full schema.
"""
if not isinstance(rng, dict):
raise TypeError(f"@tool_spec(name={name!r}): 'rng' must be a dict, got {type(rng).__name__}.")
if "uses_rng" not in rng or not isinstance(rng["uses_rng"], bool):
raise ValueError(f"@tool_spec(name={name!r}): 'rng' must have a boolean 'uses_rng' key.")
allowed_keys = {"uses_rng", "seed_param", "default_seed", "note"}
extra = set(rng) - allowed_keys
if extra:
raise ValueError(
f"@tool_spec(name={name!r}): unknown keys in 'rng': {sorted(extra)}. Allowed: {sorted(allowed_keys)}."
)
if not rng["uses_rng"]:
# Deterministic tools shouldn't carry a seed_param / default_seed.
if rng.get("seed_param") is not None or rng.get("default_seed") is not None:
raise ValueError(
f"@tool_spec(name={name!r}): 'seed_param' / 'default_seed' must be omitted when uses_rng is False."
)
return
seed_param = rng.get("seed_param")
if seed_param is not None and not isinstance(seed_param, str):
raise TypeError(
f"@tool_spec(name={name!r}): 'seed_param' must be a string or None, got {type(seed_param).__name__}."
)
default_seed = rng.get("default_seed")
if default_seed is not None and not isinstance(default_seed, int):
raise TypeError(
f"@tool_spec(name={name!r}): 'default_seed' must be an int or None, got {type(default_seed).__name__}."
)
if seed_param is None and default_seed is not None:
raise ValueError(f"@tool_spec(name={name!r}): 'default_seed' requires a 'seed_param' name.")
# ---------------------------------------------------------------------------
# Decorator
# ---------------------------------------------------------------------------
[docs]
def tool_spec( # noqa: PLR0913
name: str,
description: str,
*,
input_model: type[BaseModel],
examples: str = "",
category: str = "",
rng: dict[str, Any] | None = None,
) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
"""Mark a function as an agent-callable tool.
Parameters
----------
name:
Unique tool name exposed to the LLM (snake_case).
description:
Natural-language description of what the tool does and when to use it.
input_model:
A :class:`pydantic.BaseModel` subclass describing the tool's
input. The model is the single source of truth for both the
function's Python signature and the MCP JSON Schema (derived
via :meth:`BaseModel.model_json_schema`). Every input model
**must** set ``model_config = ConfigDict(extra="forbid")`` so
unknown keys are rejected (closes the SEC-15
kwarg-injection vector at the schema layer).
The decorated function then receives the parsed model as its
single positional argument::
class FooInput(BaseModel):
model_config = ConfigDict(extra="forbid")
values: list[float] = Field(..., min_length=2)
@tool_spec(name="foo", description="...", input_model=FooInput)
def foo(spec: FooInput) -> dict:
return {"n": len(spec.values)}
examples:
Optional natural-language -> tool-call mappings appended to
``description`` so the LLM sees worked examples.
category:
Optional category string (e.g. ``"univariate"``). Used for
filtering with :func:`get_tool_specs`.
rng:
Optional reproducibility metadata; see :func:`_validate_rng_metadata`
for the contract.
Returns
-------
Callable
The original function, with an attached ``_tool_spec`` dict
and ``_input_model`` attribute.
"""
if rng is not None:
_validate_rng_metadata(name, rng)
if not (isinstance(input_model, type) and issubclass(input_model, BaseModel)):
raise TypeError(
f"@tool_spec(name={name!r}): input_model must be a pydantic BaseModel subclass; got {input_model!r}."
)
extra_policy = input_model.model_config.get("extra")
if extra_policy != "forbid":
raise ValueError(
f"@tool_spec(name={name!r}): input_model {input_model.__name__} "
"must set model_config = ConfigDict(extra='forbid') so unknown "
"kwargs are rejected (ENG-04 / SEC-15 contract)."
)
def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
"""Attach the assembled tool spec to ``func`` and return it unchanged."""
full_description = description
if examples:
full_description = f"{description}\n\nExamples\n--------\n{examples}"
spec: dict[str, Any] = {
"name": name,
"description": full_description,
"input_schema": input_model.model_json_schema(),
}
if category:
spec["category"] = category
if rng is not None:
spec["rng"] = dict(rng)
func._tool_spec = spec # type: ignore[attr-defined]
func._input_model = input_model # type: ignore[attr-defined]
_TOOL_REGISTRY[name] = func
return func
return decorator
# ---------------------------------------------------------------------------
# Serialisation helper
# ---------------------------------------------------------------------------
[docs]
def clean(value: Any) -> Any: # noqa: PLR0911, ANN401
"""Recursively convert numpy scalars / arrays to plain Python types.
All ``tools.py`` modules should call ``clean(result)`` before returning
so that every tool output is JSON-serialisable.
"""
if isinstance(value, dict):
# Keys need unwrapping too: json.dumps rejects numpy scalar keys (a
# dict keyed by np.int64 group labels from a pandas groupby is the
# common case), and previously they leaked through untouched. Only
# numpy scalars are unwrapped - other key types pass through so
# hashability is never at risk.
return {(k.item() if isinstance(k, np.generic) else k): clean(v) for k, v in value.items()}
if isinstance(value, (list, tuple, set, frozenset)):
return [clean(v) for v in value]
# np.bool_ subclasses neither np.integer nor Python bool, and json.dumps
# rejects it; it must be tested before the numeric branches.
if isinstance(value, np.bool_):
return bool(value)
if isinstance(value, np.integer):
return int(value)
if isinstance(value, np.floating):
v = float(value)
return None if math.isnan(v) or math.isinf(v) else v
if isinstance(value, float):
return None if math.isnan(value) or math.isinf(value) else value
if isinstance(value, np.ndarray):
return clean(value.tolist())
if isinstance(value, np.generic):
# Catch-all for the remaining numpy scalar types (datetime64,
# complex, str_, ...): unwrap to the closest Python equivalent.
return clean(value.item())
return value
# ---------------------------------------------------------------------------
# Discovery
# ---------------------------------------------------------------------------
def _import_tool_module(module: str) -> None:
"""Import a single ``tools.py`` module for discovery.
A genuinely missing module - typically an uninstalled optional third-party
dependency that the tools module imports - is tolerated but logged, so the
dropped tool category is visible rather than silent. Any other
:class:`ImportError` (for example a bad name imported inside the module) is a
real bug and is allowed to propagate rather than being silently swallowed,
which would make the whole tool category vanish without a trace.
"""
import importlib # noqa: PLC0415
try:
importlib.import_module(module)
except ModuleNotFoundError as exc:
# ModuleNotFoundError is raised for ANY unresolvable module path,
# including a typo inside the tools module itself or a renamed
# internal module after a refactor. Only tolerate it when the module
# that is actually missing (exc.name) is a third-party dependency;
# a missing first-party module is a real bug and must propagate,
# exactly as the docstring above promises.
missing = exc.name or ""
if missing.split(".")[0] == "process_improve":
raise
logger.warning("Tool module %r not loaded (missing dependency %r): %s", module, missing, exc)
[docs]
def discover_tools() -> None:
"""Import all ``tools.py`` modules to populate the tool registry.
This is called lazily on the first :func:`get_tool_specs` invocation.
It is safe to call multiple times (subsequent calls are no-ops).
"""
global _discovery_done # noqa: PLW0603
if _discovery_done:
return
for module in [
"process_improve.univariate.tools",
"process_improve.multivariate.tools",
"process_improve.monitoring.tools",
"process_improve.regression.tools",
"process_improve.bivariate.tools",
"process_improve.experiments.tools",
"process_improve.batch.tools",
"process_improve.visualization.tools",
"process_improve.simulation.tools",
"process_improve.sensory.tools",
"process_improve.recipes",
]:
_import_tool_module(module)
_discovery_done = True
# ---------------------------------------------------------------------------
# Public helpers
# ---------------------------------------------------------------------------
[docs]
def get_tool_specs(
names: list[str] | None = None,
category: str | None = None,
) -> list[dict[str, Any]]:
"""Return tool specs in the format expected by the Anthropic ``tools=`` parameter.
Parameters
----------
names:
Optional allow-list of tool names to include. When *None* (default)
all registered tools are returned.
category:
Optional category filter (e.g. ``"univariate"``). When provided, only
tools whose ``category`` matches are returned.
Returns
-------
list[dict]
Each dict has keys ``"name"``, ``"description"``, and
``"input_schema"`` as required by the Anthropic API. Tools that
opt in via ``rng=`` on the decorator also carry an ``"rng"`` key
describing their reproducibility contract; see :func:`tool_spec`.
"""
discover_tools()
registry = _TOOL_REGISTRY
if names is not None:
registry = {k: v for k, v in registry.items() if k in names}
specs = [func._tool_spec for func in registry.values()] # type: ignore[attr-defined]
if category is not None:
specs = [s for s in specs if s.get("category") == category]
return specs
[docs]
def execute_tool_call(tool_name: str, tool_input: dict[str, Any]) -> Any: # noqa: ANN401
"""Dispatch a single tool call from an Anthropic ``tool_use`` content block.
The input dict is validated via
``input_model.model_validate(tool_input)``. Unknown keys raise
``ToolInputInvalidError`` (closes the SEC-15 ``confirmed=True``
kwarg-injection at the schema layer). The parsed pydantic model is
passed to the tool function as a single positional argument.
Parameters
----------
tool_name:
The ``name`` field from the ``tool_use`` block.
tool_input:
The ``input`` dict from the ``tool_use`` block.
Returns
-------
Any
Whatever the tool function returns (typically a JSON-serialisable
``dict``).
Raises
------
ValueError
If *tool_name* is not in the registry.
ToolInputInvalidError
If the input fails pydantic validation.
"""
discover_tools()
if tool_name not in _TOOL_REGISTRY:
available = sorted(_TOOL_REGISTRY)
raise ValueError(f"Unknown tool {tool_name!r}. Available tools: {available}")
func = _TOOL_REGISTRY[tool_name]
model_cls: type[BaseModel] = func._input_model # type: ignore[attr-defined]
try:
parsed = model_cls.model_validate(tool_input)
except ValidationError as exc:
raise ToolInputInvalidError(
f"Input to tool {tool_name!r} failed validation: {exc}",
details={"tool_name": tool_name, "errors": exc.errors()},
) from exc
return func(parsed)
# ---------------------------------------------------------------------------
# Safety wrapper re-exports
# ---------------------------------------------------------------------------
# Callers that expose the registry over an untrusted transport should use
# ``safe_execute_tool_call`` from ``process_improve.tool_safety`` instead of
# ``execute_tool_call``. The names are re-exported here for discoverability.
from process_improve.tool_safety import ( # noqa: E402
ToolInputInvalidError,
ToolInputTooLargeError,
ToolMemoryExceededError,
ToolSafetyError,
ToolTimeoutError,
safe_execute_tool_call,
)
__all__ = [
"ToolInputInvalidError",
"ToolInputTooLargeError",
"ToolMemoryExceededError",
"ToolSafetyError",
"ToolTimeoutError",
"clean",
"discover_tools",
"execute_tool_call",
"get_tool_specs",
"safe_execute_tool_call",
"tool_spec",
]