From 18f6a98bd6c46e2e6c6595a0d309c8d708616928 Mon Sep 17 00:00:00 2001
From: Parth Sareen <parth.sareen@ollama.com>
Date: Wed, 11 Dec 2024 17:17:36 -0800
Subject: [PATCH] llama: enable JSON schema key ordering for generating
 grammars (#8055)

---
 llama/grammar_test.go                         |  6 +----
 llama/json-schema-to-grammar.cpp              |  2 +-
 ...ntain-ordering-for-rules-for-grammar.patch | 22 +++++++++++++++++++
 3 files changed, 24 insertions(+), 6 deletions(-)
 create mode 100644 llama/patches/0012-Maintain-ordering-for-rules-for-grammar.patch

diff --git a/llama/grammar_test.go b/llama/grammar_test.go
index b64976ae8..e4a34f0df 100644
--- a/llama/grammar_test.go
+++ b/llama/grammar_test.go
@@ -30,8 +30,6 @@ const issue7978JSONSchema = `{
 }`
 
 func TestIssue7978(t *testing.T) {
-	t.Skip("schema_to_grammar is broken; skipping until fixed")
-
 	g := SchemaToGrammar([]byte(issue7978JSONSchema))
 	if g == nil {
 		t.Fatal("failed to convert JSON schema to grammar")
@@ -54,8 +52,6 @@ func TestIssue7978(t *testing.T) {
 }
 
 func TestSchemaToGrammer(t *testing.T) {
-	t.Skip("schema_to_grammar is broken; skipping until fixed")
-
 	cases := []struct {
 		schema string
 		prefix []byte // nil is check as nil
@@ -63,7 +59,7 @@ func TestSchemaToGrammer(t *testing.T) {
 		{`invalid`, nil},
 
 		// Simple heuristic/smoke test
-		{`{"type":"object"}`, []byte("object ::=")},
+		{`{"type":"object"}`, []byte("root ::= object")},
 	}
 
 	for _, c := range cases {
diff --git a/llama/json-schema-to-grammar.cpp b/llama/json-schema-to-grammar.cpp
index e5191967b..958244a02 100644
--- a/llama/json-schema-to-grammar.cpp
+++ b/llama/json-schema-to-grammar.cpp
@@ -417,7 +417,7 @@ class SchemaConverter {
 private:
     std::function<json(const std::string &)> _fetch_json;
     bool _dotall;
-    std::map<std::string, std::string> _rules;
+    std::unordered_map<std::string, std::string> _rules;
     std::unordered_map<std::string, json> _refs;
     std::unordered_set<std::string> _refs_being_resolved;
     std::vector<std::string> _errors;
diff --git a/llama/patches/0012-Maintain-ordering-for-rules-for-grammar.patch b/llama/patches/0012-Maintain-ordering-for-rules-for-grammar.patch
new file mode 100644
index 000000000..33b504ec1
--- /dev/null
+++ b/llama/patches/0012-Maintain-ordering-for-rules-for-grammar.patch
@@ -0,0 +1,22 @@
+From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
+From: ParthSareen <parth.sareen@ollama.com>
+Date: Wed, 11 Dec 2024 15:37:32 -0800
+Subject: [PATCH] Maintain ordering for rules for grammar
+
+---
+ common/json-schema-to-grammar.cpp | 2 +-
+ 1 file changed, 1 insertion(+), 1 deletion(-)
+
+diff --git a/common/json-schema-to-grammar.cpp b/common/json-schema-to-grammar.cpp
+index dadc18c8..2a8dbd22 100644
+--- a/common/json-schema-to-grammar.cpp
++++ b/common/json-schema-to-grammar.cpp
+@@ -391,7 +391,7 @@ class SchemaConverter {
+ private:
+     std::function<json(const std::string &)> _fetch_json;
+     bool _dotall;
+-    std::map<std::string, std::string> _rules;
++    std::unordered_map<std::string, std::string> _rules;
+     std::unordered_map<std::string, json> _refs;
+     std::unordered_set<std::string> _refs_being_resolved;
+     std::vector<std::string> _errors;