From 4aa5dac768a5255667097db1c043581196fa46b2 Mon Sep 17 00:00:00 2001
From: Ettore Di Giacinto <mudler@users.noreply.github.com>
Date: Thu, 3 Aug 2023 23:32:30 +0200
Subject: [PATCH] feat: update integer, number and string rules - allow
 primitives as root types (#862)

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
---
 pkg/grammar/json_schema.go      | 14 ++++++++++----
 pkg/grammar/json_schema_test.go |  5 ++++-
 2 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/pkg/grammar/json_schema.go b/pkg/grammar/json_schema.go
index 8c9b7008..5a8ebd0e 100644
--- a/pkg/grammar/json_schema.go
+++ b/pkg/grammar/json_schema.go
@@ -15,10 +15,13 @@ var (
 
 	PRIMITIVE_RULES = map[string]string{
 		"boolean": `("true" | "false") space`,
-		"number":  `[0-9]+ space`,                    // TODO complete
-		"integer": `[0-9]+ space`,                    // TODO complete
-		"string":  `"\"" [ \t!#-\[\]-~]* "\"" space`, // TODO complete
-		"null":    `"null" space`,
+		"number":  `("-"? ([0-9] | [1-9] [0-9]*)) ("." [0-9]+)? ([eE] [-+]? [0-9]+)? space`,
+		"integer": `("-"? ([0-9] | [1-9] [0-9]*)) space`,
+		"string": `"\"" (
+			[^"\\] |
+			"\\" (["\\/bfnrt] | "u" [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F])
+		  )* "\"" space`,
+		"null": `"null" space`,
 	}
 
 	INVALID_RULE_CHARS_RE     = regexp.MustCompile(`[^a-zA-Z0-9-]+`)
@@ -176,6 +179,9 @@ func (sc *JSONSchemaConverter) visit(schema map[string]interface{}, name string,
 		if !exists {
 			panic(fmt.Sprintf("Unrecognized schema: %v", schema))
 		}
+		if ruleName == "root" {
+			schemaType = "root"
+		}
 		return sc.addRule(schemaType, primitiveRule)
 	}
 }
diff --git a/pkg/grammar/json_schema_test.go b/pkg/grammar/json_schema_test.go
index 0d8dd992..9d4086cb 100644
--- a/pkg/grammar/json_schema_test.go
+++ b/pkg/grammar/json_schema_test.go
@@ -48,7 +48,10 @@ root ::= root-0 | root-1
 space ::= " "?
 root-0-arguments ::= "{" space "\"date\"" space ":" space string "," space "\"time\"" space ":" space string "," space "\"title\"" space ":" space string "}" space
 root-1 ::= "{" space "\"arguments\"" space ":" space root-1-arguments "," space "\"function\"" space ":" space root-1-function "}" space
-string ::= "\"" [ \t!#-\[\]-~]* "\"" space
+string ::= "\"" (
+	[^"\\] |
+	"\\" (["\\/bfnrt] | "u" [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F])
+)* "\"" space
 root-1-function ::= "\"search\""`
 )