Teknium
1545afb892
Zero behavior change; tool schemas byte-identical; golden corpus (376 guard commands, 93 heredoc forms, sanitizer/unrename cases) identical old vs new. registry.py (1263 -> 924): _memo_check per-pass check_fn memo (2 sites), _grouped/_toolset_entries toolset grouping (6 sites), _unique_env replaces _extend_unique, _attr accessor for get_schema/get_toolset_for_tool/get_emoji/ get_max_result_size, fold cache-prune loops, flatten _callable_module walk, merge guard try-blocks, docstring/comment compaction (all WHY kept). schema_sanitizer.py (511 -> 392): _rewrite bottom-up tree map shared by _strip_ref_siblings/strip_nullable_unions/collapse_const_unions, _dict_nodes generator replaces nested _walk closure, collapsed _const_branch_type guards. self_repo_guard.py (678 -> 517): _scope_keys state machine as one if/elif ladder, heredoc opener parsed by one regex (_HEREDOC_OPENER_RE), _masked_line inlined, _operator_before via rstrip, folded tail expressions. plugin_guard.py (235 -> 163): positional Finding ctor, packed tables, docs. project_tools.py (181 -> 156): _activated shared by create/switch, _ACTIONS dict dispatch replaces the if-chain in _handle_project. path_security.py (24 -> 18): unused logger/logging import dropped.
393 lines
18 KiB
Python
393 lines
18 KiB
Python
"""Sanitize tool JSON schemas for broad LLM-backend compatibility.
|
|
|
|
Strict backends reject shapes OpenAI/Anthropic accept: llama.cpp's grammar converter fails
|
|
on ``{"type": "object"}`` without ``properties``, bare-string schemas and ``type`` arrays;
|
|
Anthropic rejects nullable ``anyOf`` at the top of ``input_schema``; Fireworks rejects
|
|
``default`` beside ``$ref``; Codex rejects top-level combinators. This module walks the
|
|
final schema tree on a deep copy and fixes only those shapes.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import copy
|
|
import logging
|
|
import re
|
|
from typing import Any, Callable
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# Anthropic (and Bedrock/Vertex/Azure fronting it) reject property keys not matching this;
|
|
# one bad key anywhere in the tools array 400s the request (Cloudflare's MCP ships 61).
|
|
_PROP_KEY_RE = re.compile(r"^[a-zA-Z0-9_.-]{1,64}$")
|
|
_PROP_KEY_BAD_CHARS = re.compile(r"[^a-zA-Z0-9_.-]")
|
|
|
|
_UNION_KEYS = ("anyOf", "oneOf")
|
|
# Outer-node metadata carried onto a union's replacement node.
|
|
_UNION_META_KEYS = ("title", "description", "default", "examples")
|
|
|
|
|
|
def _empty_object() -> dict:
|
|
return {"type": "object", "properties": {}}
|
|
|
|
|
|
def _rewrite(schema: Any, fn: Callable[[dict], Any]) -> Any:
|
|
"""Bottom-up map over a schema tree: lists/dicts recurse, then *fn* sees each dict."""
|
|
if isinstance(schema, list):
|
|
return [_rewrite(item, fn) for item in schema]
|
|
if not isinstance(schema, dict):
|
|
return schema
|
|
return fn({k: _rewrite(v, fn) for k, v in schema.items()})
|
|
|
|
|
|
def sanitize_property_key(key: str) -> str:
|
|
"""Deterministically map an arbitrary property key to a conforming one."""
|
|
return _PROP_KEY_BAD_CHARS.sub("_", key)[:64] or "param"
|
|
|
|
|
|
def _rename_property_keys(props: dict, path: str) -> dict[str, str]:
|
|
"""{original_key: conforming_key} for one properties dict (identity entries omitted).
|
|
Deterministic (insertion order, numeric suffixes on collision) so the model-visible
|
|
schema and the dispatch-time reverse map from the registry's original schema agree."""
|
|
renames: dict[str, str] = {}
|
|
taken = {k for k in props if _PROP_KEY_RE.match(k)}
|
|
for key in props:
|
|
if _PROP_KEY_RE.match(key):
|
|
continue
|
|
base = sanitize_property_key(key)
|
|
candidate, i = base, 2
|
|
while candidate in taken:
|
|
suffix = f"_{i}"
|
|
candidate = base[: 64 - len(suffix)] + suffix
|
|
i += 1
|
|
taken.add(candidate)
|
|
renames[key] = candidate
|
|
logger.debug(
|
|
"schema_sanitizer[%s]: renamed property key %r -> %r "
|
|
"(provider key-pattern compat)", path, key, candidate)
|
|
return renames
|
|
|
|
|
|
def unrename_tool_args(params_schema: Any, args: Any) -> Any:
|
|
"""Map sanitized property keys in model-emitted args back to wire names. ``params_schema``
|
|
is the ORIGINAL registry schema; recurses into objects/array items; unknown keys pass."""
|
|
if not isinstance(params_schema, dict) or not isinstance(args, dict):
|
|
return args
|
|
props = params_schema.get("properties")
|
|
if not isinstance(props, dict):
|
|
return args
|
|
reverse = {v: k for k, v in _rename_property_keys(props, "<unrename>").items()}
|
|
out = {}
|
|
for key, value in args.items():
|
|
orig = reverse.get(key, key)
|
|
subschema = props.get(orig)
|
|
if isinstance(subschema, dict):
|
|
if isinstance(value, dict):
|
|
value = unrename_tool_args(subschema, value)
|
|
elif isinstance(value, list) and isinstance(subschema.get("items"), dict):
|
|
value = [
|
|
unrename_tool_args(subschema["items"], item) if isinstance(item, dict) else item
|
|
for item in value]
|
|
out[orig] = value
|
|
return out
|
|
|
|
|
|
def sanitize_tool_schemas(tools: list[dict]) -> list[dict]:
|
|
"""Deep-copied ``tools`` (OpenAI format) with sanitized parameter schemas; safe to mutate."""
|
|
if not tools:
|
|
return tools
|
|
return [_sanitize_single_tool(tool) for tool in tools]
|
|
|
|
|
|
def _sanitize_single_tool(tool: dict) -> dict:
|
|
out = copy.deepcopy(tool)
|
|
fn = out.get("function") if isinstance(out, dict) else None
|
|
if not isinstance(fn, dict):
|
|
return out
|
|
params = fn.get("parameters")
|
|
if not isinstance(params, dict): # missing / non-dict → minimal valid shape
|
|
fn["parameters"] = _empty_object()
|
|
return out
|
|
name = fn.get("name", "<tool>")
|
|
top = _sanitize_node(params, path=name)
|
|
# Guarantee the top level is an object with properties.
|
|
if not isinstance(top, dict):
|
|
top = {}
|
|
top["type"] = "object"
|
|
if not isinstance(top.get("properties"), dict):
|
|
top["properties"] = {}
|
|
# The recursive pass only handles array-form ``type: [X, "null"]``; collapse anyOf unions
|
|
# here, keeping ``nullable: true`` so ``model_tools._schema_allows_null`` still coerces.
|
|
top = strip_nullable_unions(top, keep_nullable_hint=True)
|
|
top = _strip_top_level_combinators(top, path=name)
|
|
fn["parameters"] = _strip_ref_siblings(top)
|
|
return out
|
|
|
|
|
|
# Sibling keywords strict JSON Schema validators reject alongside ``$ref``.
|
|
_REF_FORBIDDEN_SIBLINGS = frozenset({"default"})
|
|
|
|
|
|
def _strip_ref_siblings(node: Any) -> Any:
|
|
"""Recursively drop forbidden siblings of ``$ref`` (Fireworks rejects ``default`` there)."""
|
|
def strip(out: dict) -> dict:
|
|
if "$ref" in out:
|
|
for key in _REF_FORBIDDEN_SIBLINGS:
|
|
out.pop(key, None)
|
|
return out
|
|
return _rewrite(node, strip)
|
|
|
|
|
|
_TOP_LEVEL_FORBIDDEN_KEYS = ("allOf", "anyOf", "oneOf", "enum", "not")
|
|
|
|
|
|
def _strip_top_level_combinators(params: dict, *, path: str = "<tool>") -> dict:
|
|
"""Drop combinators from the TOP level only (Codex rejects them there). They are usually
|
|
conditional-required hints, so validity is unchanged (handlers re-validate); nested
|
|
combinators are preserved."""
|
|
if not isinstance(params, dict):
|
|
return params
|
|
out = dict(params)
|
|
for key in _TOP_LEVEL_FORBIDDEN_KEYS:
|
|
if key in out:
|
|
logger.debug(
|
|
"schema_sanitizer[%s]: stripped top-level %r combinator "
|
|
"from tool parameters (strict-backend compat)",
|
|
path, key)
|
|
out.pop(key, None)
|
|
return out
|
|
|
|
|
|
def _is_null_branch(item: Any) -> bool:
|
|
return isinstance(item, dict) and item.get("type") == "null"
|
|
|
|
|
|
def _carry_union_meta(outer: dict, replacement: dict, *, skip_default_on_ref: bool) -> None:
|
|
"""Copy outer-union metadata onto *replacement* where absent."""
|
|
for meta_key in _UNION_META_KEYS:
|
|
if meta_key in outer and meta_key not in replacement:
|
|
# ``default`` is illegal alongside ``$ref`` on strict backends.
|
|
if skip_default_on_ref and meta_key == "default" and "$ref" in replacement:
|
|
continue
|
|
replacement[meta_key] = outer[meta_key]
|
|
|
|
|
|
def strip_nullable_unions(schema: Any, *, keep_nullable_hint: bool = True) -> Any:
|
|
"""Collapse ``anyOf``/``oneOf`` nullable unions to the single non-null branch. MCP/Pydantic
|
|
optional fields arrive as ``{"anyOf": [{"type": "string"}, {"type": "null"}]}``; Anthropic
|
|
rejects the null branch and optionality is already in the parent's ``required``. Only
|
|
collapses when a null branch was dropped AND exactly one non-null branch survives.
|
|
``keep_nullable_hint`` sets ``nullable: true`` for runtime ``"null"`` → ``None`` coercion."""
|
|
def collapse(stripped: dict) -> Any:
|
|
for key in _UNION_KEYS:
|
|
variants = stripped.get(key)
|
|
if not isinstance(variants, list):
|
|
continue
|
|
non_null = [item for item in variants if not _is_null_branch(item)]
|
|
if len(non_null) == 1 and len(non_null) != len(variants):
|
|
replacement = dict(non_null[0]) if isinstance(non_null[0], dict) else {}
|
|
if keep_nullable_hint:
|
|
replacement.setdefault("nullable", True)
|
|
_carry_union_meta(stripped, replacement, skip_default_on_ref=True)
|
|
return _rewrite(replacement, collapse)
|
|
return stripped
|
|
return _rewrite(schema, collapse)
|
|
|
|
|
|
_CONST_PRIMITIVE_TYPES: dict[type, str] = {
|
|
bool: "boolean", int: "integer", float: "number", str: "string"}
|
|
|
|
|
|
def _const_branch_type(branch: Any) -> str | None:
|
|
"""JSON-Schema primitive type of a pure ``const`` branch, else None: a primitive ``const``
|
|
whose declared ``type`` (if any) matches; only ``title``/``description`` may accompany it."""
|
|
if not isinstance(branch, dict) or "const" not in branch \
|
|
or set(branch) - {"const", "type", "title", "description"}:
|
|
return None
|
|
# ``type(value)`` lookup (not isinstance): bool is a subclass of int.
|
|
json_type = _CONST_PRIMITIVE_TYPES.get(type(branch["const"]))
|
|
return json_type if json_type is not None and branch.get("type") in (None, json_type) else None
|
|
|
|
|
|
def collapse_const_unions(schema: Any) -> Any:
|
|
"""Collapse ``anyOf``/``oneOf`` unions of same-typed consts to ``enum`` (ported from
|
|
block/goose ``tool_schema_normalize.rs``, Apache-2.0). Rust/TS MCP servers emit
|
|
``{"anyOf": [{"const": "red"}, {"const": "green"}]}``, which strict backends mishandle.
|
|
Applies only when EVERY non-null branch is a pure ``const`` of one primitive type
|
|
(``bool`` never merges with ``integer``); one ``{"type": "null"}`` branch is tolerated
|
|
and recorded as ``nullable: true``. Branch order is kept; outer metadata carried over;
|
|
input never mutated."""
|
|
def collapse(out: dict) -> Any:
|
|
for key in _UNION_KEYS:
|
|
variants = out.get(key)
|
|
if not isinstance(variants, list) or not variants:
|
|
continue
|
|
null_branches = [i for i in variants if _is_null_branch(i) and "const" not in i]
|
|
const_branches = [item for item in variants if item not in null_branches]
|
|
if len(null_branches) > 1 or not const_branches:
|
|
continue
|
|
branch_types = {_const_branch_type(item) for item in const_branches}
|
|
if len(branch_types) != 1 or None in branch_types:
|
|
continue
|
|
replacement: dict = {
|
|
"type": branch_types.pop(), "enum": [item["const"] for item in const_branches]}
|
|
if null_branches:
|
|
replacement["nullable"] = True
|
|
_carry_union_meta(out, replacement, skip_default_on_ref=False)
|
|
return replacement
|
|
return out
|
|
return _rewrite(schema, collapse)
|
|
|
|
|
|
_BARE_TYPE_NAMES = frozenset({"object", "string", "number", "integer", "boolean", "array", "null"})
|
|
# Keys whose values are NOT schemas (recursing would treat "path" as a bare-string schema).
|
|
_NON_SCHEMA_LIST_KEYS = frozenset({"required", "enum", "examples", "dependentRequired"})
|
|
|
|
|
|
def _normalize_type_array(value: list, out: dict) -> None:
|
|
"""Normalize a ``type: [...]`` array into *out* (llama.cpp and Gemini-via-OpenAI reject
|
|
arrays). Per AI-SDK: one non-null type → ``type: X`` (+ ``nullable`` if ``null`` present);
|
|
several → ``anyOf`` of single-type schemas so EVERY branch survives; none → ``null`` or
|
|
the object fallback. Ported from anomalyco/opencode#31877."""
|
|
has_null = "null" in value
|
|
non_null = [t for t in value if isinstance(t, str) and t != "null"]
|
|
if len(non_null) == 1:
|
|
out["type"] = non_null[0]
|
|
elif len(non_null) >= 2:
|
|
out["anyOf"] = [{"type": t} for t in non_null]
|
|
else:
|
|
out["type"] = "null" if has_null else "object"
|
|
return
|
|
if has_null:
|
|
out.setdefault("nullable", True)
|
|
|
|
|
|
def _sanitize_node(node: Any, path: str) -> Any:
|
|
"""Recursively sanitize a JSON-Schema fragment: bare-string schemas become ``{"type":
|
|
<value>}`` (unknown strings → permissive object); object nodes gain ``properties: {}``;
|
|
``type`` arrays are normalized; property keys are renamed to the provider-safe pattern
|
|
and ``required`` follows, with entries missing from ``properties`` pruned."""
|
|
if isinstance(node, str):
|
|
if node in _BARE_TYPE_NAMES:
|
|
logger.debug(
|
|
"schema_sanitizer[%s]: replacing bare-string schema %r with {'type': %r}",
|
|
path, node, node)
|
|
return _empty_object() if node == "object" else {"type": node}
|
|
logger.debug(
|
|
"schema_sanitizer[%s]: replacing non-schema string %r "
|
|
"with empty object schema", path, node)
|
|
return _empty_object()
|
|
if isinstance(node, list):
|
|
return [_sanitize_node(item, f"{path}[{i}]") for i, item in enumerate(node)]
|
|
if not isinstance(node, dict):
|
|
return node
|
|
|
|
# Renames computed up front so ``required`` remaps even when it precedes ``properties``.
|
|
prop_renames: dict[str, str] = {}
|
|
if isinstance(node.get("properties"), dict):
|
|
prop_renames = _rename_property_keys(node["properties"], f"{path}.properties")
|
|
out: dict = {}
|
|
for key, value in node.items():
|
|
if key == "type" and isinstance(value, list):
|
|
_normalize_type_array(value, out)
|
|
elif key in {"properties", "$defs", "definitions"} and isinstance(value, dict):
|
|
renames = prop_renames if key == "properties" else {}
|
|
out[key] = {
|
|
renames.get(sub_k, sub_k): _sanitize_node(sub_v, f"{path}.{key}.{renames.get(sub_k, sub_k)}")
|
|
for sub_k, sub_v in value.items()}
|
|
elif key in {"items", "additionalProperties"}:
|
|
# Bool ``additionalProperties`` is valid; bool ``items`` is non-standard but preserved.
|
|
out[key] = value if isinstance(value, bool) else _sanitize_node(value, f"{path}.{key}")
|
|
elif key in {"anyOf", "oneOf", "allOf"} and isinstance(value, list):
|
|
out[key] = [_sanitize_node(item, f"{path}.{key}[{i}]") for i, item in enumerate(value)]
|
|
elif key in _NON_SCHEMA_LIST_KEYS:
|
|
if key == "required" and prop_renames and isinstance(value, list):
|
|
out[key] = [prop_renames.get(r, r) if isinstance(r, str) else r for r in value]
|
|
else:
|
|
out[key] = copy.deepcopy(value) if isinstance(value, (list, dict)) else value
|
|
else:
|
|
out[key] = _sanitize_node(value, f"{path}.{key}") if isinstance(value, (dict, list)) else value
|
|
if out.get("type") == "object":
|
|
if not isinstance(out.get("properties"), dict):
|
|
out["properties"] = {}
|
|
if isinstance(out.get("required"), list):
|
|
props = out.get("properties") or {}
|
|
valid = [r for r in out["required"] if isinstance(r, str) and r in props]
|
|
if not valid:
|
|
out.pop("required", None)
|
|
elif len(valid) != len(out["required"]):
|
|
out["required"] = valid
|
|
return out
|
|
|
|
|
|
# ---- Reactive strips — only invoked after a backend rejects a schema -----------------------
|
|
_STRIP_ON_RECOVERY_KEYS = frozenset({"pattern", "format"})
|
|
|
|
|
|
def _dict_nodes(node: Any):
|
|
"""Pre-order walk yielding every dict node; each is yielded before its values are visited,
|
|
so a consumer may mutate it in place."""
|
|
if isinstance(node, dict):
|
|
yield node
|
|
for v in node.values():
|
|
yield from _dict_nodes(v)
|
|
elif isinstance(node, list):
|
|
for item in node:
|
|
yield from _dict_nodes(item)
|
|
|
|
|
|
def _reactive_strip(
|
|
tools: list[dict], strip_node: Callable[[dict], int], log_msg: str) -> tuple[list[dict], int]:
|
|
"""Apply *strip_node* (returns keywords removed) to every dict node of each tool's
|
|
parameters, in place. Handles OpenAI (``{"function": {"parameters": ..}}``) and Responses
|
|
(``{"name": .., "parameters": ..}``) formats. Returns ``(tools, stripped_count)``."""
|
|
if not tools:
|
|
return tools, 0
|
|
stripped = 0
|
|
for tool in tools:
|
|
if not isinstance(tool, dict):
|
|
continue
|
|
fn = tool.get("function")
|
|
params = fn.get("parameters") if isinstance(fn, dict) else None
|
|
if not isinstance(params, dict):
|
|
params = tool.get("parameters")
|
|
if isinstance(params, dict):
|
|
stripped += sum(strip_node(node) for node in _dict_nodes(params))
|
|
if stripped:
|
|
logger.info(log_msg, stripped)
|
|
return tools, stripped
|
|
|
|
|
|
def strip_pattern_and_format(tools: list[dict]) -> tuple[list[dict], int]:
|
|
"""Strip ``pattern``/``format`` keywords from tool schemas, in place. Reactive: only after
|
|
llama.cpp's grammar converter rejected a schema (its regex engine is a small ECMAScript
|
|
subset), since cloud providers use these as prompting hints. Only strips beside
|
|
``type``/combinators, so a property literally *named* ``pattern`` is untouched."""
|
|
def _strip(node: dict) -> int:
|
|
if not ("type" in node or "anyOf" in node or "oneOf" in node or "allOf" in node):
|
|
return 0
|
|
hits = [k for k in node if k in _STRIP_ON_RECOVERY_KEYS]
|
|
for k in hits:
|
|
node.pop(k, None)
|
|
return len(hits)
|
|
return _reactive_strip(
|
|
tools, _strip,
|
|
"schema_sanitizer: stripped %d pattern/format keyword(s) from "
|
|
"tool schemas (llama.cpp grammar-parse recovery)")
|
|
|
|
|
|
def strip_slash_enum(tools: list[dict]) -> tuple[list[dict], int]:
|
|
"""Strip ``enum`` keywords whose string values contain ``/``, in place. xAI compiles
|
|
schemas to a grammar that rejects ``/`` in enum values (HTTP 400 before any token) —
|
|
typically MCP enums of HuggingFace model IDs. The constraint is a prompting hint only."""
|
|
def _strip(node: dict) -> int:
|
|
enum_val = node.get("enum")
|
|
if isinstance(enum_val, list) and any(isinstance(v, str) and "/" in v for v in enum_val):
|
|
node.pop("enum", None)
|
|
return 1
|
|
return 0
|
|
return _reactive_strip(
|
|
tools, _strip,
|
|
"schema_sanitizer: stripped %d enum keyword(s) containing '/' "
|
|
"from tool schemas (xAI Responses grammar-compile recovery)")
|