Resolve tool schema $refs before MLX chat template rendering#177
Open
Adi2K wants to merge 1 commit into
Open
Conversation
Adi2K
force-pushed
the
fix/normalize-tool-schema-types
branch
from
July 21, 2026 16:18
aa87b06 to
920d933
Compare
$refs before MLX chat template rendering
GenerationSchema hoists named nested schemas into $defs and emits the
referencing property as a bare {"$ref": "#/$defs/<name>"} with no "type".
withResolvedRoot() resolves only the top-level root reference, so a
nested-object or union tool parameter reaches the model as a typeless stub.
Cloud backends accept this, because they dereference $ref server-side.
On-device chat templates do not: Gemma's format_parameters applies
`value['type'] | upper` to every property and throws "upper filter requires
string" on the stub's missing type, failing the whole tool-enabled request.
convertToolToMLXSpec now resolves references before the template runs:
- resolveToolSchemaRefs inlines each $ref with a resolved copy of its
definition, so a nested object regains its {"type": "object", ...} shape.
Draft-07 "definitions" tables resolve the same way, since Pydantic v1 and
many MCP servers emit those. Keys beside a $ref merge onto the target.
- Every subschema slot is walked, because Dict[str, Model] emits its $ref
under additionalProperties and Tuple[Model, int] under prefixItems. The
tables are dropped whether or not a slot was walked, so a $ref outliving
them reaches the template dangling: any the walk leaves collapses too.
- Expansion is bounded, since these schemas can come from third-party MCP
servers. Cycles and pointers this resolver cannot follow collapse to
{"type": "string"}. A definition referenced twice by one parent is a DAG
rather than a cycle, so a path-only guard re-inlines every shared subtree:
a 2 KB 22-level diamond expands to ~21M nodes. A node budget caps that,
and a depth cap keeps a long chain or deep literal schema off the stack.
- normalizeToolSchemaTypes then gives any residual typeless node a
renderable scalar type, coerces non-string "type" values, and replaces a
"properties" or "items" that is not a schema mapping.
Well-formed schemas pass through unchanged, and cloud backends are untouched.
Adi2K
force-pushed
the
fix/normalize-tool-schema-types
branch
from
July 21, 2026 16:43
920d933 to
b8b27ee
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Any tool with a nested-object parameter currently fails every Gemma request with
runtime("upper filter requires string"), before a token is generated.GenerationSchemaencodes a nested tool parameter as a bare$refinto$defs, andwithResolvedRoot()resolves only the top-level root reference, so the nested part reaches the chat template as a typeless stub. For a toolsearch_records(filter: { q: String }):Cloud backends accept
$refbecause they dereference server-side; on-device templates dereference nothing, so the MLX tool path has to resolve before rendering.convertToolToMLXSpecnow runsnormalizeToolSchemaTypes(resolveToolSchemaRefs(_:)). Draft-07definitionsresolve alongside$defs, since Pydantic v1 and many MCP servers emit those.Expansion is bounded, because these schemas can come from third-party MCP servers: cycles, unfollowable pointers, and anything past the depth (64) or node (5,000) cap collapse to
{"type": "string"}rather than throwing. The node cap is the load-bearing one — a definition referenced twice by one parent is a DAG rather than a cycle, so a cycle guard alone still re-inlines every shared subtree, and a 2 KB 22-level diamond expands to ~21M nodes. llama.cpp hits the same wall and sidesteps it by emitting a named grammar rule per$ref, which isn't available here since a template needs a literal schema tree.Resolution walks every subschema slot (including
additionalProperties,prefixItems,patternPropertiesand tuple-formitems, which Pydantic emits forDict[str, Model]andTuple[...]fields), and a residual sweep guarantees no$refsurvives anywhere else — dropping the definition table while leaving a pointer behind would be worse than not resolving at all. Slots carrying instance data rather than subschemas (const,default,enum,examples) pass through untouched, since a$refkey there is a literal value.Well-formed schemas pass through byte-identical, and cloud backends are untouched.
Two things I'd value your opinion on:
{"type": "string"}, which asserts a type that may not be true. A permissive{}would be more honest, but templates that indexvalue['type']need the key present.