fix(router): fix openapi and add jsonschema validation

docs/openapi.json

@@ -1022,6 +1022,57 @@
           }
         }
       },
+      "GrammarType": {
+        "oneOf": [
+          {
+            "type": "object",
+            "required": [
+              "type",
+              "value"
+            ],
+            "properties": {
+              "type": {
+                "type": "string",
+                "enum": [
+                  "json"
+                ]
+              },
+              "value": {
+                "type": "string",
+                "description": "A string that represents a [JSON Schema](https://json-schema.org/).\n\nJSON Schema is a declarative language that allows to annotate JSON documents\nwith types and descriptions.",
+                "example": {
+                  "properties": {
+                    "location": {
+                      "type": "string"
+                    }
+                  }
+                }
+              }
+            }
+          },
+          {
+            "type": "object",
+            "required": [
+              "type",
+              "value"
+            ],
+            "properties": {
+              "type": {
+                "type": "string",
+                "enum": [
+                  "regex"
+                ]
+              },
+              "value": {
+                "type": "string"
+              }
+            }
+          }
+        ],
+        "discriminator": {
+          "propertyName": "type"
+        }
+      },
       "Info": {
         "type": "object",
         "required": [

router/Cargo.toml

@@ -22,6 +22,7 @@ text-generation-client = { path = "client" }
 clap = { version = "4.4.5", features = ["derive", "env"] }
 futures = "0.3.28"
 hf-hub = { version = "0.3.0", features = ["tokio"] }
+jsonschema = "0.17.1"
 metrics = "0.21.1"
 metrics-exporter-prometheus = { version = "0.12.1", features = [] }
 nohash-hasher = "0.2.0"

router/src/lib.rs

@@ -46,6 +46,7 @@ impl HubTokenizerConfig {
 }
 
 mod json_object_or_string_to_string {
+    use jsonschema::JSONSchema;
     use serde::{Deserialize, Deserializer};
     use serde_json::Value;
 
@@ -57,6 +58,10 @@ mod json_object_or_string_to_string {
     {
         let value = Value::deserialize(deserializer)?;
 
+        JSONSchema::options()
+            .compile(&value)
+            .map_err(|e| serde::de::Error::custom(format!("invalid JSONSchema: {e}")))?;
+
         match value {
             Value::String(s) => Ok(s),
             // Safely handle serialization and return an error if it fails
@@ -70,13 +75,18 @@ mod json_object_or_string_to_string {
     }
 }
 
-#[derive(Clone, Debug, Deserialize)]
+#[derive(Clone, Debug, Deserialize, ToSchema)]
 #[serde(tag = "type", content = "value")]
 pub(crate) enum GrammarType {
+    /// A string that represents a [JSON Schema](https://json-schema.org/).
+    ///
+    /// JSON Schema is a declarative language that allows to annotate JSON documents
+    /// with types and descriptions.
     #[serde(
         rename = "json",
         deserialize_with = "json_object_or_string_to_string::deserialize"
     )]
+    #[schema(example = json ! ({"properties": {"location":{"type": "string"}}}))]
     Json(String),
     #[serde(rename = "regex")]
     Regex(String),

router/src/server.rs

@@ -5,7 +5,7 @@ use crate::validation::ValidationError;
 use crate::{
     BestOfSequence, ChatCompletion, ChatCompletionChoice, ChatCompletionChunk, ChatCompletionDelta,
     ChatCompletionLogprobs, ChatRequest, CompatGenerateRequest, Details, ErrorResponse,
-    FinishReason, GenerateParameters, GenerateRequest, GenerateResponse, HubModelInfo,
+    FinishReason, GenerateParameters, GenerateRequest, GenerateResponse, GrammarType, HubModelInfo,
     HubTokenizerConfig, Infer, Info, Message, PrefillToken, SimpleToken, StreamDetails,
     StreamResponse, Token, TokenizeResponse, Validation,
 };
@@ -800,6 +800,7 @@ pub async fn run(
     Info,
     CompatGenerateRequest,
     GenerateRequest,
+    GrammarType,
     ChatRequest,
     Message,
     ChatCompletionChoice,