"""(c) Kevin Dunn, 2010-2026. MIT License.
The ``@tool_spec`` decorator, the registry it fills, and ``clean``.
These are the parts of :mod:`process_improve.tool_spec` that every ``tools.py``
module imports. They live apart from tool discovery, which names every
``tools.py`` module, so that a tool module depends on this file alone and not,
through discovery, on every other tool module. ``tool_spec.py`` re-exports all
of it, and callers import it from there.
"""
from __future__ import annotations
import math
from collections.abc import Callable
from typing import Any
import numpy as np
from pydantic import BaseModel
# ---------------------------------------------------------------------------
# Registry
# ---------------------------------------------------------------------------
#: Maps tool name -> decorated callable. Populated by ``@tool_spec``.
_TOOL_REGISTRY: dict[str, Callable[..., Any]] = {}
# ---------------------------------------------------------------------------
# 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
# ---------------------------------------------------------------------------
# ---------------------------------------------------------------------------
# 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
