pass schema in test

predibase · Feb 9, 2024 · df88b16 · df88b16
1 parent 5781073
commit df88b16
Showing 1 changed file with 20 additions and 0 deletions.
diff --git a/server/tests/conftest.py b/server/tests/conftest.py
@@ -5,13 +5,33 @@
 
 @pytest.fixture
 def default_pb_parameters():
+    schema = """{
+    "$defs": {
+        "Armor": {
+            "enum": ["leather", "chainmail", "plate"],
+            "title": "Armor",
+            "type": "string"
+        }
+    },
+    "properties": {
+        "name": {"maxLength": 10, "title": "Name", "type": "string"},
+        "age": {"title": "Age", "type": "integer"},
+        "armor": {"$ref": "#/$defs/Armor"},
+        "strength": {"title": "Strength", "type": "integer"}\
+    },
+    "required": ["name", "age", "armor", "strength"],
+    "title": "Character",
+    "type": "object"
+}"""
+
     return generate_pb2.NextTokenChooserParameters(
         temperature=1.0,
         repetition_penalty=1.0,
         top_k=0,
         top_p=1.0,
         typical_p=1.0,
         do_sample=False,
+        schema=schema,
     )