From 61676fb5066fe42f4586d8f951548fcb55cd75bf Mon Sep 17 00:00:00 2001
From: Jeffrey Morgan <jmorganca@gmail.com>
Date: Tue, 14 Jan 2025 12:55:45 -0800
Subject: [PATCH] llama: move grammar tests to llama_test.go (#8411)

---
 llama/grammar/grammar_test.go | 107 ----------------------------------
 llama/llama_test.go           | 104 +++++++++++++++++++++++++++++++++
 2 files changed, 104 insertions(+), 107 deletions(-)
 delete mode 100644 llama/grammar/grammar_test.go

diff --git a/llama/grammar/grammar_test.go b/llama/grammar/grammar_test.go
deleted file mode 100644
index 373652ab4..000000000
--- a/llama/grammar/grammar_test.go
+++ /dev/null
@@ -1,107 +0,0 @@
-package grammar
-
-import (
-	"bufio"
-	"bytes"
-	"strings"
-	"testing"
-
-	"github.com/ollama/ollama/llama"
-)
-
-// https://github.com/ollama/ollama/issues/7978
-const issue7978JSONSchema = `{
-  "type": "object",
-  "properties": {
-    "steps": {
-      "type": "array",
-      "items": {
-        "type": "object",
-        "properties": {
-          "explanation": { "type": "string" },
-          "output": { "type": "string" },
-          "nested": {
-            "type": "object",
-            "properties": {
-              "deep": { "type": "string" }
-            }
-          }
-        },
-        "required": ["explanation", "output"],
-        "additionalProperties": false
-      }
-    },
-    "final_answer": { "type": "string" },
-    "01_numbered_key": { "type": "string" },
-    "numbers": {
-      "type": "array",
-      "items": { "type": "number" }
-    },
-    "booleans": {
-      "type": "array", 
-      "items": { "type": "boolean" }
-    },
-    "mixed": {
-      "type": "array",
-      "items": {
-        "oneOf": [
-          { "type": "string" },
-          { "type": "number" },
-          { "type": "boolean" }
-        ]
-      }
-    }
-  },
-  "required": ["steps", "final_answer"],
-  "additionalProperties": false
-}`
-
-func TestIssue7978(t *testing.T) {
-	g := llama.SchemaToGrammar([]byte(issue7978JSONSchema))
-	if g == nil {
-		t.Fatal("failed to convert JSON schema to grammar")
-	}
-
-	t.Logf("grammar:\n%s", g)
-	t.Log()
-
-	var got string
-	s := bufio.NewScanner(bytes.NewReader(g))
-	for s.Scan() {
-		line := strings.TrimSpace(s.Text())
-		step, _, _ := strings.Cut(line, " ::= ")
-		step = strings.TrimSpace(step)
-		if step == "root" {
-			got = line
-		}
-	}
-
-	want := `root ::= "{" space steps-kv "," space final-answer-kv ( "," space ( 01-numbered-key-kv 01-numbered-key-rest | numbers-kv numbers-rest | booleans-kv booleans-rest | mixed-kv ) )? "}" space`
-	if got != want {
-		t.Errorf("root =\n%qwant:\n%q", got, want)
-	}
-}
-
-func TestSchemaToGrammer(t *testing.T) {
-	cases := []struct {
-		schema string
-		prefix []byte // nil is check as nil
-	}{
-		{`invalid`, nil},
-
-		// Simple heuristic/smoke test
-		{`{"type":"object"}`, []byte("root ::= object")},
-	}
-
-	for _, c := range cases {
-		t.Run("x", func(t *testing.T) {
-			g := llama.SchemaToGrammar([]byte(c.schema))
-			if c.prefix == nil && g != nil {
-				t.Fatalf("grammar = %v, want nil", g)
-			}
-			if !bytes.HasPrefix(g, c.prefix) {
-				t.Errorf("grammar = %q, want %q", g, c.prefix)
-			}
-		})
-	}
-}
diff --git a/llama/llama_test.go b/llama/llama_test.go
index 5f835d683..b550d1d86 100644
--- a/llama/llama_test.go
+++ b/llama/llama_test.go
@@ -1 +1,105 @@
 package llama
+
+import (
+	"bufio"
+	"bytes"
+	"strings"
+	"testing"
+)
+
+// https://github.com/ollama/ollama/issues/7978
+const issue7978JSONSchema = `{
+  "type": "object",
+  "properties": {
+    "steps": {
+      "type": "array",
+      "items": {
+        "type": "object",
+        "properties": {
+          "explanation": { "type": "string" },
+          "output": { "type": "string" },
+          "nested": {
+            "type": "object",
+            "properties": {
+              "deep": { "type": "string" }
+            }
+          }
+        },
+        "required": ["explanation", "output"],
+        "additionalProperties": false
+      }
+    },
+    "final_answer": { "type": "string" },
+    "01_numbered_key": { "type": "string" },
+    "numbers": {
+      "type": "array",
+      "items": { "type": "number" }
+    },
+    "booleans": {
+      "type": "array", 
+      "items": { "type": "boolean" }
+    },
+    "mixed": {
+      "type": "array",
+      "items": {
+        "oneOf": [
+          { "type": "string" },
+          { "type": "number" },
+          { "type": "boolean" }
+        ]
+      }
+    }
+  },
+  "required": ["steps", "final_answer"],
+  "additionalProperties": false
+}`
+
+func TestIssue7978(t *testing.T) {
+	g := SchemaToGrammar([]byte(issue7978JSONSchema))
+	if g == nil {
+		t.Fatal("failed to convert JSON schema to grammar")
+	}
+
+	t.Logf("grammar:\n%s", g)
+	t.Log()
+
+	var got string
+	s := bufio.NewScanner(bytes.NewReader(g))
+	for s.Scan() {
+		line := strings.TrimSpace(s.Text())
+		step, _, _ := strings.Cut(line, " ::= ")
+		step = strings.TrimSpace(step)
+		if step == "root" {
+			got = line
+		}
+	}
+
+	want := `root ::= "{" space steps-kv "," space final-answer-kv ( "," space ( 01-numbered-key-kv 01-numbered-key-rest | numbers-kv numbers-rest | booleans-kv booleans-rest | mixed-kv ) )? "}" space`
+	if got != want {
+		t.Errorf("root =\n%qwant:\n%q", got, want)
+	}
+}
+
+func TestSchemaToGrammer(t *testing.T) {
+	cases := []struct {
+		schema string
+		prefix []byte // nil is check as nil
+	}{
+		{`invalid`, nil},
+
+		// Simple heuristic/smoke test
+		{`{"type":"object"}`, []byte("root ::= object")},
+	}
+
+	for _, c := range cases {
+		t.Run("x", func(t *testing.T) {
+			g := SchemaToGrammar([]byte(c.schema))
+			if c.prefix == nil && g != nil {
+				t.Fatalf("grammar = %v, want nil", g)
+			}
+			if !bytes.HasPrefix(g, c.prefix) {
+				t.Errorf("grammar = %q, want %q", g, c.prefix)
+			}
+		})
+	}
+}