Fix JSON schema generation to accept schemas without 'title' property

dottxt-ai · Oct 3, 2023 · 81c4c12 · 81c4c12
1 parent 4ce76ff
commit 81c4c12
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git a/outlines/text/json_schema.py b/outlines/text/json_schema.py
@@ -176,7 +176,7 @@ def expand_json_schema(
                 expanded_properties[name] = value
 
         return {
-            "title": raw_schema["title"],
+            **({"title": raw_schema["title"]} if "title" in raw_schema else {}),
             "type": raw_schema["type"],
             "properties": expanded_properties,
         }

diff --git a/tests/text/test_json_schema.py b/tests/text/test_json_schema.py
@@ -189,6 +189,18 @@ def test_json_schema():
     ]
 
 
+def test_json_schema_no_titles():
+    schema = '{"type": "object", "properties": {"user_id": {"type": "integer"}, "name": {"type": "string"}}, "required": ["user_id", "name"]}'
+    schedule = build_schedule_from_schema(schema)
+    assert schedule == [
+        '\\{[\\n ]*"user_id"[\\n ]*:[\\n ]*',
+        {"type": "integer"},
+        '[\\n ]*,[\\n ]*"name"[\\n ]*:[\\n ]*',
+        {"type": "string"},
+        "[\\n ]*\\}",
+    ]
+
+
 def test_json_schema_with_property_ref():
     schema = """{
         "title": "User",