# Copyright 2025 Google LLC. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Tests for the schema module. Note: This file contains test helper classes that intentionally have few public methods. The too-few-public-methods warnings are expected. """ import dataclasses from unittest import mock import warnings from absl.testing import absltest from absl.testing import parameterized from langextract import schema as lx_schema from langextract.core import base_model from langextract.core import data from langextract.core import exceptions from langextract.core import format_handler as fh from langextract.core import schema from langextract.providers import schemas def _openai_extraction_items(openai_schema): return openai_schema.schema_dict["properties"][data.EXTRACTIONS_KEY]["items"] def _openai_variant(openai_schema, extraction_class): for variant in _openai_extraction_items(openai_schema)["anyOf"]: if extraction_class in variant["properties"]: return variant raise AssertionError(f"Missing OpenAI schema variant for {extraction_class}") def _openai_attribute_properties(openai_schema, extraction_class): variant = _openai_variant(openai_schema, extraction_class) attributes_key = f"{extraction_class}{data.ATTRIBUTE_SUFFIX}" return variant["properties"][attributes_key]["anyOf"][0]["properties"] class BaseSchemaTest(absltest.TestCase): """Tests for BaseSchema abstract class.""" def test_abstract_methods_required(self): """Test that BaseSchema cannot be instantiated directly.""" with self.assertRaises(TypeError): schema.BaseSchema() # pylint: disable=abstract-class-instantiated def test_subclass_must_implement_all_methods(self): """Test that subclasses must implement all abstract methods.""" class IncompleteSchema(schema.BaseSchema): # pylint: disable=too-few-public-methods @classmethod def from_examples(cls, examples_data, attribute_suffix="_attributes"): return cls() with self.assertRaises(TypeError): IncompleteSchema() # pylint: disable=abstract-class-instantiated class BaseLanguageModelSchemaTest(absltest.TestCase): """Tests for BaseLanguageModel schema methods.""" def test_get_schema_class_returns_none_by_default(self): """Test that get_schema_class returns None by default.""" class TestModel(base_model.BaseLanguageModel): # pylint: disable=too-few-public-methods def infer(self, batch_prompts, **kwargs): yield [] self.assertIsNone(TestModel.get_schema_class()) def test_apply_schema_stores_instance(self): """Test that apply_schema stores the schema instance.""" class TestModel(base_model.BaseLanguageModel): # pylint: disable=too-few-public-methods def infer(self, batch_prompts, **kwargs): yield [] model = TestModel() mock_schema = mock.Mock(spec=schema.BaseSchema) model.apply_schema(mock_schema) self.assertEqual(model._schema, mock_schema) model.apply_schema(None) self.assertIsNone(model._schema) class GeminiSchemaTest(parameterized.TestCase): @parameterized.named_parameters( dict( testcase_name="empty_extractions", examples_data=[], expected_schema={ "type": "object", "properties": { data.EXTRACTIONS_KEY: { "type": "array", "items": { "type": "object", "properties": {}, }, }, }, "required": [data.EXTRACTIONS_KEY], }, ), dict( testcase_name="single_extraction_no_attributes", examples_data=[ data.ExampleData( text="Patient has diabetes.", extractions=[ data.Extraction( extraction_text="diabetes", extraction_class="condition", ) ], ) ], expected_schema={ "type": "object", "properties": { data.EXTRACTIONS_KEY: { "type": "array", "items": { "type": "object", "properties": { "condition": {"type": "string"}, "condition_attributes": { "type": "object", "properties": { "_unused": {"type": "string"}, }, "nullable": True, }, }, }, }, }, "required": [data.EXTRACTIONS_KEY], }, ), dict( testcase_name="single_extraction", examples_data=[ data.ExampleData( text="Patient has diabetes.", extractions=[ data.Extraction( extraction_text="diabetes", extraction_class="condition", attributes={"chronicity": "chronic"}, ) ], ) ], expected_schema={ "type": "object", "properties": { data.EXTRACTIONS_KEY: { "type": "array", "items": { "type": "object", "properties": { "condition": {"type": "string"}, "condition_attributes": { "type": "object", "properties": { "chronicity": {"type": "string"}, }, "nullable": True, }, }, }, }, }, "required": [data.EXTRACTIONS_KEY], }, ), dict( testcase_name="multiple_extraction_classes", examples_data=[ data.ExampleData( text="Patient has diabetes.", extractions=[ data.Extraction( extraction_text="diabetes", extraction_class="condition", attributes={"chronicity": "chronic"}, ) ], ), data.ExampleData( text="Patient is John Doe", extractions=[ data.Extraction( extraction_text="John Doe", extraction_class="patient", attributes={"id": "12345"}, ) ], ), ], expected_schema={ "type": "object", "properties": { data.EXTRACTIONS_KEY: { "type": "array", "items": { "type": "object", "properties": { "condition": {"type": "string"}, "condition_attributes": { "type": "object", "properties": { "chronicity": {"type": "string"} }, "nullable": True, }, "patient": {"type": "string"}, "patient_attributes": { "type": "object", "properties": { "id": {"type": "string"}, }, "nullable": True, }, }, }, }, }, "required": [data.EXTRACTIONS_KEY], }, ), ) def test_from_examples_constructs_expected_schema( self, examples_data, expected_schema ): gemini_schema = schemas.gemini.GeminiSchema.from_examples(examples_data) actual_schema = gemini_schema.schema_dict self.assertEqual(actual_schema, expected_schema) def test_to_provider_config_returns_response_schema(self): """Test that to_provider_config returns the correct provider kwargs.""" examples_data = [ data.ExampleData( text="Test text", extractions=[ data.Extraction( extraction_class="test_class", extraction_text="test extraction", ) ], ) ] gemini_schema = schemas.gemini.GeminiSchema.from_examples(examples_data) provider_config = gemini_schema.to_provider_config() self.assertIn("response_schema", provider_config) self.assertEqual( provider_config["response_schema"], gemini_schema.schema_dict ) def test_requires_raw_output_returns_true(self): """Test that GeminiSchema requires raw output.""" examples_data = [ data.ExampleData( text="Test text", extractions=[ data.Extraction( extraction_class="test_class", extraction_text="test extraction", ) ], ) ] gemini_schema = schemas.gemini.GeminiSchema.from_examples(examples_data) self.assertTrue(gemini_schema.requires_raw_output) class OpenAISchemaTest(parameterized.TestCase): """Tests for OpenAI structured output schema generation.""" def test_response_format_returns_json_schema_response_format(self): """OpenAI schema exposes Chat Completions structured outputs.""" examples_data = [ data.ExampleData( text="Patient has diabetes.", extractions=[ data.Extraction( extraction_text="diabetes", extraction_class="condition", attributes={"chronicity": "chronic"}, ) ], ) ] openai_schema = schemas.openai.OpenAISchema.from_examples(examples_data) response_format = openai_schema.response_format self.assertEqual( response_format, { "type": "json_schema", "json_schema": { "name": "langextract_extractions", "schema": openai_schema.schema_dict, "strict": True, }, }, ) self.assertIsNot( response_format["json_schema"]["schema"], openai_schema.schema_dict ) def test_to_provider_config_uses_provider_schema_hook(self): """OpenAI schema state is applied after provider construction.""" openai_schema = schemas.openai.OpenAISchema.from_examples([]) provider_config = openai_schema.to_provider_config() self.assertEmpty(provider_config) def test_from_examples_constructs_strict_openai_schema(self): """OpenAI schema uses strict-compatible extraction variants.""" examples_data = [ data.ExampleData( text="Patient has diabetes.", extractions=[ data.Extraction( extraction_text="diabetes", extraction_class="condition", attributes={"chronicity": "chronic"}, ), data.Extraction( extraction_text="metformin", extraction_class="medication", attributes={"route": "oral"}, ), ], ) ] openai_schema = schemas.openai.OpenAISchema.from_examples(examples_data) self.assertEqual( openai_schema.schema_dict, { "type": "object", "properties": { data.EXTRACTIONS_KEY: { "type": "array", "items": { "anyOf": [ { "type": "object", "properties": { "condition": {"type": "string"}, "condition_attributes": { "anyOf": [ { "type": "object", "properties": { "chronicity": { "anyOf": [ {"type": "string"}, {"type": "null"}, ] } }, "required": ["chronicity"], "additionalProperties": False, }, {"type": "null"}, ] }, }, "required": [ "condition", "condition_attributes", ], "additionalProperties": False, }, { "type": "object", "properties": { "medication": {"type": "string"}, "medication_attributes": { "anyOf": [ { "type": "object", "properties": { "route": { "anyOf": [ {"type": "string"}, {"type": "null"}, ] } }, "required": ["route"], "additionalProperties": False, }, {"type": "null"}, ] }, }, "required": [ "medication", "medication_attributes", ], "additionalProperties": False, }, ] }, } }, "required": [data.EXTRACTIONS_KEY], "additionalProperties": False, }, ) def test_from_examples_preserves_list_attribute_schema(self): """OpenAI schema accepts list attributes from examples.""" examples_data = [ data.ExampleData( text="Patient has diabetes with fatigue.", extractions=[ data.Extraction( extraction_text="diabetes", extraction_class="condition", attributes={"symptoms": ["fatigue"]}, ) ], ) ] openai_schema = schemas.openai.OpenAISchema.from_examples(examples_data) self.assertEqual( _openai_attribute_properties(openai_schema, "condition")["symptoms"], { "anyOf": [ {"type": "array", "items": {"type": "string"}}, {"type": "null"}, ] }, ) def test_from_examples_empty_examples_allow_empty_extraction_objects(self): """OpenAI schema handles empty example sets deterministically.""" openai_schema = schemas.openai.OpenAISchema.from_examples([]) self.assertEqual( _openai_extraction_items(openai_schema), { "type": "object", "properties": {}, "required": [], "additionalProperties": False, }, ) def test_validate_format_rejects_yaml(self): """OpenAI structured outputs are JSON-only.""" openai_schema = schemas.openai.OpenAISchema.from_examples([]) format_handler = fh.FormatHandler(format_type=data.FormatType.YAML) with self.assertRaisesRegex( exceptions.InferenceConfigError, "OpenAI structured output only supports JSON format", ): openai_schema.validate_format(format_handler) def test_requires_raw_output_returns_true(self): """OpenAI structured outputs emit raw JSON without fences.""" openai_schema = schemas.openai.OpenAISchema.from_examples([]) self.assertTrue(openai_schema.requires_raw_output) def test_validate_format_warns_when_fences_enabled(self): """OpenAI schema warns when raw JSON would be wrapped in fences.""" openai_schema = schemas.openai.OpenAISchema.from_examples([]) format_handler = fh.FormatHandler( format_type=data.FormatType.JSON, use_fences=True, ) with self.assertWarnsRegex( UserWarning, "OpenAI structured outputs emit native JSON" ): openai_schema.validate_format(format_handler) def test_validate_format_warns_with_wrong_wrapper_key(self): """OpenAI schema warns when resolver wrapper settings drift.""" openai_schema = schemas.openai.OpenAISchema.from_examples([]) format_handler = fh.FormatHandler( format_type=data.FormatType.JSON, use_fences=False, wrapper_key="items", ) with self.assertWarnsRegex( UserWarning, f"response_format schema expects wrapper_key='{data.EXTRACTIONS_KEY}'", ): openai_schema.validate_format(format_handler) def test_from_examples_preserves_scalar_attribute_types(self): """Scalar attribute types map to their JSON-Schema equivalents. Regression test: prior to this, every non-list attribute was coerced to a string-only union, which forced OpenAI strict mode to return scalars as strings even when examples used numbers/bools. """ examples_data = [ data.ExampleData( text="Aspirin 81 mg, daily, OTC.", extractions=[ data.Extraction( extraction_text="aspirin", extraction_class="medication", attributes={ "dose_mg": 81, "doses_per_day": 1.0, "otc": True, "route": "oral", }, ) ], ) ] openai_schema = schemas.openai.OpenAISchema.from_examples(examples_data) self.assertEqual( _openai_attribute_properties(openai_schema, "medication"), { "dose_mg": {"anyOf": [{"type": "integer"}, {"type": "null"}]}, "doses_per_day": {"anyOf": [{"type": "number"}, {"type": "null"}]}, "otc": {"anyOf": [{"type": "boolean"}, {"type": "null"}]}, "route": {"anyOf": [{"type": "string"}, {"type": "null"}]}, }, ) def test_from_examples_preserves_mixed_numeric_attribute_types(self): """Mixed numeric-like examples keep each observed JSON type.""" examples_data = [ data.ExampleData( text="Medication flag.", extractions=[ data.Extraction( extraction_text="flag", extraction_class="medication", attributes={"dose_or_flag": True}, ) ], ), data.ExampleData( text="Medication count.", extractions=[ data.Extraction( extraction_text="count", extraction_class="medication", attributes={"dose_or_flag": 1}, ) ], ), data.ExampleData( text="Medication dose.", extractions=[ data.Extraction( extraction_text="dose", extraction_class="medication", attributes={"dose_or_flag": 1.5}, ) ], ), ] openai_schema = schemas.openai.OpenAISchema.from_examples(examples_data) self.assertEqual( _openai_attribute_properties(openai_schema, "medication")[ "dose_or_flag" ], { "anyOf": [ {"type": "boolean"}, {"type": "integer"}, {"type": "number"}, {"type": "null"}, ] }, ) def test_from_examples_allows_none_attribute_values(self): """None-valued example attributes keep the strict-mode null branch.""" examples_data = [ data.ExampleData( text="Medication status is unspecified.", extractions=[ data.Extraction( extraction_text="Medication", extraction_class="medication", attributes={"status": None}, ) ], ) ] openai_schema = schemas.openai.OpenAISchema.from_examples(examples_data) self.assertEqual( _openai_attribute_properties(openai_schema, "medication")["status"], {"anyOf": [{"type": "string"}, {"type": "null"}]}, ) def test_from_examples_strict_false_emits_non_strict_response_format(self): """The strict kwarg threads through to response_format.""" openai_schema = schemas.openai.OpenAISchema.from_examples([], strict=False) self.assertFalse(openai_schema.response_format["json_schema"]["strict"]) def test_response_format_returns_isolated_schema_dict(self): """response_format callers cannot mutate the provider's schema.""" openai_schema = schemas.openai.OpenAISchema.from_examples([]) response_format = openai_schema.response_format response_format["json_schema"]["schema"]["required"].append("extra") self.assertEqual( openai_schema.schema_dict["required"], [data.EXTRACTIONS_KEY] ) def test_instance_is_frozen_and_dict_is_isolated(self): """Frozen contract + deep-copy isolate the schema from caller mutation.""" source = { "type": "object", "properties": {"x": {"type": "string"}}, "required": ["x"], "additionalProperties": False, } openai_schema = schemas.openai.OpenAISchema(schema_dict=source) with self.assertRaises(dataclasses.FrozenInstanceError): openai_schema.schema_dict = {} # pylint: disable=attribute-defined-outside-init source["properties"]["x"]["type"] = "integer" self.assertEqual( openai_schema.schema_dict["properties"]["x"], {"type": "string"} ) class SchemaValidationTest(parameterized.TestCase): """Tests for schema format validation.""" def _create_test_schema(self): """Helper to create a test schema.""" examples = [ data.ExampleData( text="Test", extractions=[ data.Extraction( extraction_class="entity", extraction_text="test", ) ], ) ] return schemas.gemini.GeminiSchema.from_examples(examples) @parameterized.named_parameters( dict( testcase_name="warns_about_fences", use_fences=True, use_wrapper=True, wrapper_key=data.EXTRACTIONS_KEY, expected_warning="fence_output=True may cause parsing issues", ), dict( testcase_name="warns_about_wrong_wrapper_key", use_fences=False, use_wrapper=True, wrapper_key="wrong_key", expected_warning="response_schema expects wrapper_key='extractions'", ), dict( testcase_name="no_warning_with_correct_settings", use_fences=False, use_wrapper=True, wrapper_key=data.EXTRACTIONS_KEY, expected_warning=None, ), ) def test_gemini_validation( self, use_fences, use_wrapper, wrapper_key, expected_warning ): """Test GeminiSchema validation with various settings.""" schema_obj = self._create_test_schema() format_handler = fh.FormatHandler( format_type=data.FormatType.JSON, use_fences=use_fences, use_wrapper=use_wrapper, wrapper_key=wrapper_key, ) with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") schema_obj.validate_format(format_handler) if expected_warning: self.assertLen( w, 1, f"Expected exactly one warning containing '{expected_warning}'", ) self.assertIn( expected_warning, str(w[0].message), f"Warning message should contain '{expected_warning}'", ) else: self.assertEmpty(w, "No warnings should be issued for correct settings") def test_base_schema_no_validation(self): """Test that base schema has no validation by default.""" schema_obj = schema.FormatModeSchema() format_handler = fh.FormatHandler( format_type=data.FormatType.JSON, use_fences=True, ) with warnings.catch_warnings(record=True) as w: warnings.simplefilter("always") schema_obj.validate_format(format_handler) self.assertEmpty( w, "FormatModeSchema should not issue validation warnings" ) class OutputSchemaHelperTest(parameterized.TestCase): """Tests for the public lx.schema output-schema builders.""" def test_extraction_item_schema_without_attributes(self): item = lx_schema.extraction_item_schema("condition") self.assertEqual( item, { "type": "object", "properties": {"condition": {"type": "string"}}, "required": ["condition"], "additionalProperties": False, }, ) def test_extraction_item_schema_with_attributes(self): item = lx_schema.extraction_item_schema( "condition", attributes={"status": {"type": "string", "enum": ["active"]}}, ) self.assertEqual( item["properties"]["condition_attributes"], { "type": "object", "properties": {"status": {"type": "string", "enum": ["active"]}}, "required": ["status"], "additionalProperties": False, }, ) self.assertCountEqual( item["required"], ["condition", "condition_attributes"] ) def test_extractions_schema_wraps_item_schema(self): item = lx_schema.extraction_item_schema("condition") envelope = lx_schema.extractions_schema(item) self.assertEqual( envelope["properties"][data.EXTRACTIONS_KEY]["items"], item ) self.assertEqual(envelope["required"], [data.EXTRACTIONS_KEY]) self.assertIs(envelope["additionalProperties"], False) lx_schema.validate_output_schema(envelope) def test_extractions_schema_wraps_multiple_items_in_any_of(self): condition = lx_schema.extraction_item_schema("condition") medication = lx_schema.extraction_item_schema("medication") envelope = lx_schema.extractions_schema(condition, medication) self.assertEqual( envelope["properties"][data.EXTRACTIONS_KEY]["items"], {"anyOf": [condition, medication]}, ) lx_schema.validate_output_schema(envelope) def test_builders_copy_input_schemas(self): attribute_schema = {"type": "string"} item = lx_schema.extraction_item_schema( "condition", attributes={"status": attribute_schema} ) attribute_schema["enum"] = ["mutated"] status_schema = item["properties"]["condition_attributes"]["properties"][ "status" ] self.assertNotIn("enum", status_schema) @parameterized.named_parameters( dict(testcase_name="empty_class", extraction_class=""), dict(testcase_name="reserved_suffix", extraction_class="foo_attributes"), dict(testcase_name="reserved_key", extraction_class="extraction_text"), ) def test_extraction_item_schema_rejects_invalid_class(self, extraction_class): with self.assertRaises(exceptions.InferenceConfigError): lx_schema.extraction_item_schema(extraction_class) class OutputSchemaValidationTest(parameterized.TestCase): """Tests for validate_output_schema envelope checks.""" def test_accepts_envelope_and_returns_isolated_copy(self): envelope = lx_schema.extractions_schema( lx_schema.extraction_item_schema("condition") ) validated = lx_schema.validate_output_schema(envelope) self.assertEqual(validated, envelope) validated["properties"]["extra"] = {"type": "string"} self.assertNotIn("extra", envelope["properties"]) @parameterized.named_parameters( dict( testcase_name="not_a_mapping", output_schema=["extractions"], error_regex="must be a mapping", ), dict( testcase_name="empty", output_schema={}, error_regex="must not be empty", ), dict( testcase_name="non_object_root", output_schema={"type": "array"}, error_regex="top-level type", ), dict( testcase_name="missing_required_extractions", output_schema={ "type": "object", "properties": { "extractions": { "type": "array", "items": { "type": "object", "properties": {"condition": {"type": "string"}}, }, } }, }, error_regex="required must include 'extractions'", ), dict( testcase_name="extractions_not_array", output_schema={ "type": "object", "required": ["extractions"], "properties": {"extractions": {"type": "object"}}, }, error_regex="array property", ), dict( testcase_name="items_not_object_schema", output_schema={ "type": "object", "required": ["extractions"], "properties": { "extractions": { "type": "array", "items": {"type": "string"}, } }, }, error_regex="inline object schema", ), dict( testcase_name="items_without_properties", output_schema={ "type": "object", "required": ["extractions"], "properties": { "extractions": { "type": "array", "items": {"type": "object"}, } }, }, error_regex="extraction-class properties", ), dict( testcase_name="reserved_item_keys", output_schema={ "type": "object", "required": ["extractions"], "properties": { "extractions": { "type": "array", "items": { "type": "object", "properties": { "extraction_class": {"type": "string"}, "extraction_text": {"type": "string"}, }, }, } }, }, error_regex="extraction_class, extraction_text", ), ) def test_rejects_invalid_envelopes(self, output_schema, error_regex): with self.assertRaisesRegex(exceptions.InferenceConfigError, error_regex): lx_schema.validate_output_schema(output_schema) class FromSchemaDictTest(absltest.TestCase): """Tests for provider from_schema_dict implementations.""" def setUp(self): super().setUp() self.envelope = lx_schema.extractions_schema( lx_schema.extraction_item_schema( "condition", attributes={ "status": {"type": "string", "enum": ["active", "resolved"]} }, ) ) def test_base_schema_rejects_user_schemas_by_default(self): with self.assertRaises(NotImplementedError): schema.FormatModeSchema.from_schema_dict(self.envelope) def test_gemini_from_schema_dict_targets_json_schema_field(self): gemini_schema = schemas.gemini.GeminiSchema.from_schema_dict(self.envelope) provider_config = gemini_schema.to_provider_config() self.assertEqual(provider_config["response_json_schema"], self.envelope) self.assertEqual(provider_config["response_mime_type"], "application/json") self.assertTrue(gemini_schema.from_output_schema) self.assertTrue(gemini_schema.requires_raw_output) def test_gemini_from_schema_dict_validates_envelope(self): with self.assertRaisesRegex( exceptions.InferenceConfigError, "array property" ): schemas.gemini.GeminiSchema.from_schema_dict({ "type": "object", "required": ["extractions"], "properties": {"extractions": {"type": "object"}}, }) def test_openai_from_schema_dict_builds_response_format(self): openai_schema = schemas.openai.OpenAISchema.from_schema_dict(self.envelope) response_format = openai_schema.response_format self.assertEqual(response_format["type"], "json_schema") self.assertEqual(response_format["json_schema"]["schema"], self.envelope) self.assertIs(response_format["json_schema"]["strict"], True) self.assertTrue(openai_schema.from_output_schema) self.assertTrue(openai_schema.requires_raw_output) def test_openai_from_schema_dict_rejects_invalid_schema_name(self): with self.assertRaisesRegex(exceptions.InferenceConfigError, "schema_name"): schemas.openai.OpenAISchema.from_schema_dict( self.envelope, schema_name="bad name!" ) if __name__ == "__main__": absltest.main()