Validate schema on data_schema setter (#46)

run-llama · Jan 22, 2025 · fcb8d4b · fcb8d4b
1 parent 5d6f8ed
commit fcb8d4b
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/llama_extract/extract.py b/llama_extract/extract.py
@@ -113,7 +113,12 @@ def data_schema(self, data_schema: SchemaInput) -> None:
             raise ValueError(
                 "data_schema must be either a dictionary or a Pydantic model"
             )
-        self._data_schema = processed_schema
+        validated_schema = self._run_in_thread(
+            self._client.llama_extract.validate_extraction_schema(
+                data_schema=processed_schema
+            )
+        )
+        self._data_schema = validated_schema.data_schema
 
     @property
     def config(self) -> ExtractConfig: