jambo/schema_converter.py

@@ -1,4 +1,4 @@
-from jambo.parser import ObjectTypeParser
+from jambo.parser import ObjectTypeParser, RefTypeParser
 from jambo.types.json_schema_type import JSONSchema
 
 from jsonschema.exceptions import SchemaError
@@ -32,18 +32,22 @@ class SchemaConverter:
         if "title" not in schema:
             raise ValueError("JSON Schema must have a title.")
 
-        if (schema_type := schema.get("type", "undefined")) != "object":
-            raise TypeError(
-                f"Invalid JSON Schema: {schema_type}. Only 'object' can be converted to Pydantic models."
-            )
+        schema_type = SchemaConverter._get_schema_type(schema)
 
-        parsed_model = ObjectTypeParser.to_model(
-            schema["title"],
-            schema.get("properties"),
-            schema.get("required"),
-            context=schema,
-            ref_cache=dict(),
-        )
+        parsed_model = None
+        match schema_type:
+            case "object":
+                parsed_model = SchemaConverter._from_object(schema)
+            case "$ref":
+                parsed_model, _ = RefTypeParser().from_properties(
+                    schema["title"],
+                    schema,
+                    context=schema,
+                    ref_cache=dict(),
+                    required=True,
+                )
+            case _:
+                raise TypeError(f"Unsupported schema type: {schema_type}")
 
         if not issubclass(parsed_model, BaseModel):
             raise TypeError(
@@ -51,3 +55,34 @@ class SchemaConverter:
             )
 
         return parsed_model
+
+    @staticmethod
+    def _from_object(schema: JSONSchema) -> type[BaseModel]:
+        """
+        Converts a JSON Schema object to a Pydantic model.
+        :param schema: The JSON Schema object to convert.
+        :return: A Pydantic model class.
+        """
+
+        if "properties" not in schema:
+            raise ValueError("JSON Schema object must have properties defined.")
+
+        return ObjectTypeParser.to_model(
+            schema["title"],
+            schema["properties"],
+            schema.get("required", []),
+            context=schema,
+            ref_cache=dict(),
+        )
+
+    @staticmethod
+    def _get_schema_type(schema: JSONSchema) -> str:
+        """
+        Returns the type of the schema.
+        :param schema: The JSON Schema to check.
+        :return: The type of the schema.
+        """
+        if "$ref" in schema:
+            return "$ref"
+
+        return schema.get("type", "undefined")

tests/test_schema_converter.py

@@ -546,6 +546,19 @@ class TestSchemaConverter(TestCase):
             },
         }
 
-        with self.assertRaises(TypeError):
-            # This should raise TypeError because the root schema is not an object
-            SchemaConverter.build(schema)
+        model = SchemaConverter.build(schema)
+
+        obj = model(
+            name="John",
+            age=30,
+            emergency_contact=model(
+                name="Jane",
+                age=28,
+            ),
+        )
+
+        self.assertEqual(obj.name, "John")
+        self.assertEqual(obj.age, 30)
+        self.assertIsInstance(obj.emergency_contact, model)
+        self.assertEqual(obj.emergency_contact.name, "Jane")
+        self.assertEqual(obj.emergency_contact.age, 28)