Repository: ggml-org/llama.cpp
Pull request: grammar : support array references in json schema

fixes #16790

related:  https://github.com/ggml-org/llama.cpp/issues/16714#issuecomment-3431179568.

The JSON schema to grammar conversion does not support referencing array items. It appears zod or the MCP library may do this to reuse schemas instead of creating a separate definition.

<details>
<summary>Example</summary>

```bash
curl http://localhost:8080/v1/chat/completions -d '{
  "model": "gpt-oss-20b",
  "messages": [
    {
      "role": "user",
      "content": "Build a binary tree that matches (A (B C D) E)"
    }
  ],
  "tools": [
    {
      "type": "function",
      "function": {
        "name": "build_binary_tree",
        "description": "build a binary tree. The left/right fields may be either a string value or another tree.",
        "parameters": {
          "type": "object",
          "properties": {
            "tree": {
              "type": "object",
              "properties": {
                "left": {
                  "anyOf": [
                    {
                      "type": "string"
                    },
                    {
                      "$ref": "#/properties/tree"
                    }
                  ]
                },
                "right": {
                  "anyOf": [
                    {
                      "$ref": "#/properties/tree/properties/left/anyOf/0"
                    },
                    {
                      "$ref": "#/properties/tree"
                    }
                  ]
                }
              },
              "additionalProperties": false
            }
          },
          "required": [
            "tree"
          ],
          "additionalProperties": false,
          "$schema": "http://json-schema.org/draft-07/schema#"
        }
      }
    }
  ],
  "tool_choice": "auto"
}'
```
</details>

The example above will fail with:

```json
{
  "error": {
    "code": 500,
    "message": "JSON schema conversion failed:\nError resolving ref #/properties/tree/properties/left/anyOf/0: 0 not in [{\"type\":\"string\"},{\"$ref\":\"#/properties/tree\"}]",
    "type": "server_error"
  }
}
```

This PR adds support for referencing array items.

Since indexes are less unique than object keys, it also renames the grammar rules derived from references.

Currently, the rules are named after the last component of the reference. E.g. `#/properties/tree => tree`. That doesn't work well with indexes, so instead this PR names them as follows:

1. Extract the content after `#`.
2. Replace non-alphanumeric characters with a dash `-`.
3. Prefix with `ref`

This results in `#/properties/tree => ref-properties-tree` as the grammar rule name.

@ochafik I would like your opinion on the rule naming. I don't know if there is any additional impact I am not seeing.

Here is the same example against this PR:

```json
{
  "model": "unsloth/gpt-oss-20b",
  "choices": [
    {
      "finish_reason": "tool_calls",
      "index": 0,
      "message": {
        "role": "assistant",
        "reasoning_content": "We need to use the function build_binary_tree. The f...",
        "content": null,
        "tool_calls": [
          {
            "type": "function",
            "function": {
              "name": "build_binary_tree",
              "arguments": "{\"tree\":{\"left\":{\"left\":\"C\",\"right\":\"D\"},\"right\":\"E\"}}"
            },
            "id": "NHd0M0zGWTx1XG6DN2NalUqMT0b8GVb3"
          }
        ]
      }
    }
  ]
}
```

And the grammar rules generated:

```
build-binary-tree-args ::= "{" space build-binary-tree-args-tree-kv "}" space
build-binary-tree-args-tree ::= "{" space  (build-binary-tree-args-tree-left-kv build-binary-tree-args-tree-left-rest | build-binary-tree-args-tree-right-kv )? "}" space
build-binary-tree-args-tree-kv ::= "\"tree\"" space ":" space build-binary-tree-args-tree
build-binary-tree-args-tree-left ::= string | build-binary-tree-args-tree-left-1
build-binary-tree-args-tree-left-1 ::= ref-properties-tree
build-binary-tree-args-tree-left-kv ::= "\"left\"" space ":" space build-binary-tree-args-tree-left
build-binary-tree-args-tree-left-rest ::= ( "," space build-binary-tree-args-tree-right-kv )?
build-binary-tree-args-tree-right ::= build-binary-tree-args-tree-right-0 | build-binary-tree-args-tree-right-1
build-binary-tree-args-tree-right-0 ::= string
build-binary-tree-args-tree-right-1 ::= ref-properties-tree
build-binary-tree-args-tree-right-kv ::= "\"right\"" space ":" space build-binary-tree-args-tree-right
build-binary-tree-call ::= "build_binary_tree"channel " 
json"? "
" build-binary-tree-args
build-binary-tree-call0 ::= "build_binary_tree" " 
json"? "
" build-binary-tree-args
channel ::= "
" ( "commentary" | "analysis" )
char ::= [^"\\\x7F\x00-\x1F] | [\\] (["\\bfnrt] | "u" [0-9a-fA-F]{4})
recipient-in-channel ::= channel " to=functions." ( build-binary-tree-call0 )
recipient-in-role ::= "
assistant"? " to=functions." ( build-binary-tree-call )
ref-properties-tree ::= "{" space  (ref-properties-tree-left-kv ref-properties-tree-left-rest | ref-properties-tree-right-kv )? "}" space
ref-properties-tree-left ::= string | ref-properties-tree-left-1
ref-properties-tree-left-1 ::= ref-properties-tree
ref-properties-tree-left-kv ::= "\"left\"" space ":" space ref-properties-tree-left
ref-properties-tree-left-rest ::= ( "," space ref-properties-tree-right-kv )?
ref-properties-tree-right ::= ref-properties-tree-right-0 | ref-properties-tree-right-1
ref-properties-tree-right-0 ::= string
ref-properties-tree-right-1 ::= ref-properties-tree
ref-properties-tree-right-kv ::= "\"right\"" space ":" space ref-properties-tree-right
root ::= recipient-in-role | recipient-in-channel
space ::= | " " | "\n"{1,2} [ \t]{0,20}
string ::= "\"" char* "\"" space
```

Changed files:
filename: common/json-schema-to-grammar.cpp
patch: @@ -601,7 +601,10 @@ class SchemaConverter {
     }
 
     std::string _resolve_ref(const std::string & ref) {
-        std::string ref_name = ref.substr(ref.find_last_of('/') + 1);
+        auto it = ref.find('#');
+        std::string ref_fragment = it != std::string::npos ? ref.substr(it + 1) : ref;
+        static const std::regex nonalphanumeric_regex(R"([^a-zA-Z0-9-]+)");
+        std::string ref_name = "ref" + std::regex_replace(ref_fragment, nonalphanumeric_regex, "-");
         if (_rules.find(ref_name) == _rules.end() && _refs_being_resolved.find(ref) == _refs_being_resolved.end()) {
             _refs_being_resolved.insert(ref);
             json resolved = _refs[ref];
@@ -774,11 +777,24 @@ class SchemaConverter {
                         std::vector<std::string> tokens = string_split(pointer, "/");
                         for (size_t i = 1; i < tokens.size(); ++i) {
                             std::string sel = tokens[i];
-                            if (target.is_null() || !target.contains(sel)) {
+                            if (target.is_object() && target.contains(sel)) {
+                                target = target[sel];
+                            } else if (target.is_array()) {
+                                size_t sel_index;
+                                try {
+                                    sel_index = std::stoul(sel);
+                                } catch (const std::invalid_argument & e) {
+                                    sel_index = target.size();
+                                }
+                                if (sel_index >= target.size()) {
+                                    _errors.push_back("Error resolving ref " + ref + ": " + sel + " not in " + target.dump());
+                                    return;
+                                }
+                                target = target[sel_index];
+                            } else {
                                 _errors.push_back("Error resolving ref " + ref + ": " + sel + " not in " + target.dump());
                                 return;
                             }
-                            target = target[sel];
                         }
                         _refs[ref] = target;
                     }

filename: examples/json_schema_to_grammar.py
patch: @@ -371,8 +371,17 @@ def visit(n: dict):
                         raise ValueError(f'Unsupported ref {ref}')
 
                     for sel in ref.split('#')[-1].split('/')[1:]:
-                        assert target is not None and sel in target, f'Error resolving ref {ref}: {sel} not in {target}'
-                        target = target[sel]
+                        assert target is not None, f'Error resolving ref {ref}: {sel} not in {target}'
+                        if isinstance(target, list):
+                            try:
+                                sel_index = int(sel)
+                            except ValueError:
+                                raise ValueError(f'Error resolving ref {ref}: {sel} not in {target}')
+                            assert 0 <= sel_index < len(target), f'Error resolving ref {ref}: {sel} not in {target}'
+                            target = target[sel_index]
+                        else:
+                            assert sel in target, f'Error resolving ref {ref}: {sel} not in {target}'
+                            target = target[sel]
 
                     self._refs[ref] = target
                 else:
@@ -547,7 +556,8 @@ def join_seq():
 
 
     def _resolve_ref(self, ref):
-        ref_name = ref.split('/')[-1]
+        ref_fragment = ref.split('#')[-1]
+        ref_name = 'ref' + re.sub(r'[^a-zA-Z0-9-]+', '-', ref_fragment)
         if ref_name not in self._rules and ref not in self._refs_being_resolved:
             self._refs_being_resolved.add(ref)
             resolved = self._refs[ref]

filename: tests/test-json-schema-to-grammar.cpp
patch: @@ -1124,9 +1124,9 @@ static void test_all(const std::string & lang, std::function<void(const TestCase
         })""",
         R"""(
             char ::= [^"\\\x7F\x00-\x1F] | [\\] (["\\bfnrt] | "u" [0-9a-fA-F]{4})
-            foo ::= "{" space foo-a-kv "}" space
-            foo-a-kv ::= "\"a\"" space ":" space string
-            root ::= foo
+            ref-definitions-foo ::= "{" space ref-definitions-foo-a-kv "}" space
+            ref-definitions-foo-a-kv ::= "\"a\"" space ":" space string
+            root ::= ref-definitions-foo
             space ::= | " " | "\n"{1,2} [ \t]{0,20}
             string ::= "\"" char* "\"" space
         )"""
@@ -1151,20 +1151,58 @@ static void test_all(const std::string & lang, std::function<void(const TestCase
             "type": "object"
         })""",
         R"""(
-            alternative-0 ::= foo
-            alternative-1 ::= bar
-            bar ::= "{" space  (bar-b-kv )? "}" space
-            bar-b-kv ::= "\"b\"" space ":" space number
+            alternative-0 ::= ref-definitions-foo
+            alternative-1 ::= ref-definitions-bar
             decimal-part ::= [0-9]{1,16}
-            foo ::= "{" space  (foo-a-kv )? "}" space
-            foo-a-kv ::= "\"a\"" space ":" space number
             integral-part ::= [0] | [1-9] [0-9]{0,15}
             number ::= ("-"? integral-part) ("." decimal-part)? ([eE] [-+]? integral-part)? space
+            ref-definitions-bar ::= "{" space  (ref-definitions-bar-b-kv )? "}" space
+            ref-definitions-bar-b-kv ::= "\"b\"" space ":" space number
+            ref-definitions-foo ::= "{" space  (ref-definitions-foo-a-kv )? "}" space
+            ref-definitions-foo-a-kv ::= "\"a\"" space ":" space number
             root ::= alternative-0 | alternative-1
             space ::= | " " | "\n"{1,2} [ \t]{0,20}
         )"""
     });
 
+    test({
+        SUCCESS,
+        "anyOf $ref",
+        R"""({
+            "properties": {
+                "a": {
+                    "anyOf": [
+                        {"type": "string"},
+                        {"type": "number"}
+                    ]
+                },
+                "b": {
+                    "anyOf": [
+                        {"$ref": "#/properties/a/anyOf/0"},
+                        {"type": "boolean"}
+                    ]
+                }
+            },
+            "type": "object"
+        })""",
+        R"""(
+            a ::= string | number
+            a-kv ::= "\"a\"" space ":" space a
+            a-rest ::= ( "," space b-kv )?
+            b ::= b-0 | boolean
+            b-0 ::= string
+            b-kv ::= "\"b\"" space ":" space b
+            boolean ::= ("true" | "false") space
+            char ::= [^"\\\x7F\x00-\x1F] | [\\] (["\\bfnrt] | "u" [0-9a-fA-F]{4})
+            decimal-part ::= [0-9]{1,16}
+            integral-part ::= [0] | [1-9] [0-9]{0,15}
+            number ::= ("-"? integral-part) ("." decimal-part)? ([eE] [-+]? integral-part)? space
+            root ::= "{" space  (a-kv a-rest | b-kv )? "}" space
+            space ::= | " " | "\n"{1,2} [ \t]{0,20}
+            string ::= "\"" char* "\"" space
+        )"""
+    });
+
     test({
         SUCCESS,
         "mix of allOf, anyOf and $ref (similar to https://json.schemastore.org/tsconfig.json)",

filename: tools/server/public_legacy/json-schema-to-grammar.mjs
patch: @@ -345,10 +345,14 @@ export class SchemaConverter {
 
           const selectors = ref.split('#')[1].split('/').slice(1);
           for (const sel of selectors) {
-            if (!target || !(sel in target)) {
+            const selIndex = parseInt(sel, 10);
+            if (target && sel in target) {
+              target = target[sel];
+            } else if (target && selIndex in target) {
+              target = target[selIndex];
+            } else {
               throw new Error(`Error resolving ref ${ref}: ${sel} not in ${JSON.stringify(target)}`);
             }
-            target = target[sel];
           }
 
           this._refs[ref] = target;
@@ -594,7 +598,8 @@ export class SchemaConverter {
   }
 
   _resolveRef(ref) {
-    let refName = ref.split('/').pop();
+    let refFragment = ref.split('#').pop();
+    let refName = 'ref' + refFragment.replace(/[^a-zA-Z0-9-]+/g, '-');
     if (!(refName in this._rules) && !this._refsBeingResolved.has(ref)) {
       this._refsBeingResolved.add(ref);
       const resolved = this._refs[ref];