wehub-resource-sync
76d991c447
Auto Update PR / update-prs (push) Has been cancelled
CI / format-check (push) Has been cancelled
CI / test (3.10) (push) Has been cancelled
CI / test (3.11) (push) Has been cancelled
CI / test (3.12) (push) Has been cancelled
CI / live-api-tests (push) Has been cancelled
CI / plugin-integration-test (push) Has been cancelled
CI / ollama-integration-test (push) Has been cancelled
CI / test-fork-pr (push) Has been cancelled
547 lines
18 KiB
Python
547 lines
18 KiB
Python
# 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.
|
|
|
|
"""Integration tests for extract function with new schema system."""
|
|
|
|
from unittest import mock
|
|
import warnings
|
|
|
|
from absl.testing import absltest
|
|
|
|
from langextract import factory
|
|
import langextract as lx
|
|
from langextract.core import data
|
|
from langextract.core import exceptions
|
|
from langextract.providers import gemini
|
|
|
|
|
|
class ExtractSchemaIntegrationTest(absltest.TestCase):
|
|
"""Tests for extract function with schema system integration."""
|
|
|
|
def setUp(self):
|
|
"""Set up test fixtures."""
|
|
super().setUp()
|
|
self.examples = [
|
|
data.ExampleData(
|
|
text="Patient has diabetes",
|
|
extractions=[
|
|
data.Extraction(
|
|
extraction_class="condition",
|
|
extraction_text="diabetes",
|
|
attributes={"severity": "moderate"},
|
|
)
|
|
],
|
|
)
|
|
]
|
|
self.test_text = "Patient has hypertension"
|
|
|
|
@mock.patch.dict("os.environ", {"GEMINI_API_KEY": "test_key"})
|
|
def test_extract_with_gemini_uses_schema(self):
|
|
"""Test that extract with Gemini automatically uses schema."""
|
|
with mock.patch(
|
|
"langextract.providers.gemini.GeminiLanguageModel.__init__",
|
|
return_value=None,
|
|
) as mock_init:
|
|
with mock.patch(
|
|
"langextract.providers.gemini.GeminiLanguageModel.infer",
|
|
return_value=iter([[mock.Mock(output='{"extractions": []}')]]),
|
|
):
|
|
with mock.patch(
|
|
"langextract.annotation.Annotator.annotate_text",
|
|
return_value=data.AnnotatedDocument(
|
|
text=self.test_text, extractions=[]
|
|
),
|
|
):
|
|
result = lx.extract(
|
|
text_or_documents=self.test_text,
|
|
prompt_description="Extract conditions",
|
|
examples=self.examples,
|
|
model_id="gemini-3.5-flash",
|
|
use_schema_constraints=True,
|
|
fence_output=None, # Let it compute
|
|
)
|
|
|
|
# Should have been called with response_schema
|
|
call_kwargs = mock_init.call_args[1]
|
|
self.assertIn("response_schema", call_kwargs)
|
|
|
|
# Result should be an AnnotatedDocument
|
|
self.assertIsInstance(result, data.AnnotatedDocument)
|
|
|
|
@mock.patch.dict("os.environ", {"OLLAMA_BASE_URL": "http://localhost:11434"})
|
|
def test_extract_with_ollama_uses_json_mode(self):
|
|
"""Test that extract with Ollama uses JSON mode."""
|
|
with mock.patch(
|
|
"langextract.providers.ollama.OllamaLanguageModel.__init__",
|
|
return_value=None,
|
|
) as mock_init:
|
|
with mock.patch(
|
|
"langextract.providers.ollama.OllamaLanguageModel.infer",
|
|
return_value=iter([[mock.Mock(output='{"extractions": []}')]]),
|
|
):
|
|
with mock.patch(
|
|
"langextract.annotation.Annotator.annotate_text",
|
|
return_value=data.AnnotatedDocument(
|
|
text=self.test_text, extractions=[]
|
|
),
|
|
):
|
|
result = lx.extract(
|
|
text_or_documents=self.test_text,
|
|
prompt_description="Extract conditions",
|
|
examples=self.examples,
|
|
model_id="gemma2:2b",
|
|
use_schema_constraints=True,
|
|
fence_output=None, # Let it compute
|
|
)
|
|
|
|
# Should have been called with format="json"
|
|
call_kwargs = mock_init.call_args[1]
|
|
self.assertIn("format", call_kwargs)
|
|
self.assertEqual(call_kwargs["format"], "json")
|
|
|
|
# Result should be an AnnotatedDocument
|
|
self.assertIsInstance(result, data.AnnotatedDocument)
|
|
|
|
def test_extract_explicit_fence_respected(self):
|
|
"""Test that explicit fence_output is respected in extract."""
|
|
with mock.patch(
|
|
"langextract.providers.gemini.GeminiLanguageModel.__init__",
|
|
return_value=None,
|
|
):
|
|
with mock.patch(
|
|
"langextract.providers.gemini.GeminiLanguageModel.infer",
|
|
return_value=iter([[mock.Mock(output='{"extractions": []}')]]),
|
|
):
|
|
with mock.patch(
|
|
"langextract.annotation.Annotator.__init__", return_value=None
|
|
) as mock_annotator_init:
|
|
with mock.patch(
|
|
"langextract.annotation.Annotator.annotate_text",
|
|
return_value=data.AnnotatedDocument(
|
|
text=self.test_text, extractions=[]
|
|
),
|
|
):
|
|
_ = lx.extract(
|
|
text_or_documents=self.test_text,
|
|
prompt_description="Extract conditions",
|
|
examples=self.examples,
|
|
model_id="gemini-3.5-flash",
|
|
api_key="test_key",
|
|
use_schema_constraints=True,
|
|
fence_output=True, # Explicitly set
|
|
)
|
|
|
|
# Annotator should be created with format_handler that has use_fences=True
|
|
call_kwargs = mock_annotator_init.call_args[1]
|
|
self.assertIn("format_handler", call_kwargs)
|
|
self.assertTrue(call_kwargs["format_handler"].use_fences)
|
|
|
|
def test_extract_gemini_schema_deprecation_warning(self):
|
|
"""Test that passing gemini_schema triggers deprecation warning."""
|
|
with mock.patch(
|
|
"langextract.providers.gemini.GeminiLanguageModel.__init__",
|
|
return_value=None,
|
|
):
|
|
with mock.patch(
|
|
"langextract.providers.gemini.GeminiLanguageModel.infer",
|
|
return_value=iter([[mock.Mock(output='{"extractions": []}')]]),
|
|
):
|
|
with mock.patch(
|
|
"langextract.annotation.Annotator.annotate_text",
|
|
return_value=data.AnnotatedDocument(
|
|
text=self.test_text, extractions=[]
|
|
),
|
|
):
|
|
with warnings.catch_warnings(record=True) as w:
|
|
warnings.simplefilter("always")
|
|
|
|
_ = lx.extract(
|
|
text_or_documents=self.test_text,
|
|
prompt_description="Extract conditions",
|
|
examples=self.examples,
|
|
model_id="gemini-3.5-flash",
|
|
api_key="test_key",
|
|
language_model_params={
|
|
"gemini_schema": "some_schema"
|
|
}, # Deprecated
|
|
)
|
|
|
|
# Should have triggered deprecation warning
|
|
deprecation_warnings = [
|
|
warning
|
|
for warning in w
|
|
if issubclass(warning.category, FutureWarning)
|
|
and "gemini_schema" in str(warning.message)
|
|
]
|
|
self.assertGreater(len(deprecation_warnings), 0)
|
|
|
|
def test_extract_no_schema_when_disabled(self):
|
|
"""Test that no schema is used when use_schema_constraints=False."""
|
|
# Create a mock instance with required attributes
|
|
mock_model = mock.MagicMock()
|
|
mock_model._schema = None
|
|
mock_model._fence_output_override = None
|
|
mock_model.gemini_schema = None
|
|
mock_model.requires_fence_output = True
|
|
mock_model.infer.return_value = iter(
|
|
[[mock.Mock(output='{"extractions": []}')]]
|
|
)
|
|
|
|
# Track the kwargs passed to the constructor
|
|
constructor_kwargs = {}
|
|
|
|
def mock_constructor(**kwargs):
|
|
constructor_kwargs.update(kwargs)
|
|
return mock_model
|
|
|
|
with mock.patch(
|
|
"langextract.providers.gemini.GeminiLanguageModel",
|
|
side_effect=mock_constructor,
|
|
):
|
|
with mock.patch(
|
|
"langextract.annotation.Annotator.annotate_text",
|
|
return_value=data.AnnotatedDocument(
|
|
text=self.test_text, extractions=[]
|
|
),
|
|
):
|
|
_ = lx.extract(
|
|
text_or_documents=self.test_text,
|
|
prompt_description="Extract conditions",
|
|
examples=self.examples,
|
|
model_id="gemini-3.5-flash",
|
|
api_key="test_key",
|
|
use_schema_constraints=False, # Disabled
|
|
)
|
|
|
|
# Should NOT have response_schema when schema constraints are disabled
|
|
self.assertNotIn("response_schema", constructor_kwargs)
|
|
self.assertNotIn("gemini_schema", constructor_kwargs)
|
|
|
|
@mock.patch("langextract.factory.create_model")
|
|
def test_validation_triggers_warning_for_gemini(self, mock_create_model):
|
|
"""Test that Gemini schema validation triggers warnings."""
|
|
|
|
# Setup mock model with Gemini schema
|
|
mock_model = mock.MagicMock()
|
|
mock_model.requires_fence_output = True
|
|
mock_model.infer.return_value = [
|
|
[mock.MagicMock(output='{"extractions": []}', score=1.0)]
|
|
]
|
|
|
|
# Create a mock Gemini schema with validate_format that issues warnings
|
|
mock_schema = mock.MagicMock()
|
|
|
|
def mock_validate_format(format_handler, level=None):
|
|
# Simulate the warning that would be issued
|
|
warnings.warn(
|
|
"Gemini outputs native JSON via"
|
|
" response_mime_type='application/json'",
|
|
UserWarning,
|
|
stacklevel=3,
|
|
)
|
|
|
|
mock_schema.validate_format = mock_validate_format
|
|
mock_model.schema = mock_schema
|
|
|
|
mock_create_model.return_value = mock_model
|
|
|
|
# Run extraction with warnings captured
|
|
with warnings.catch_warnings(record=True) as w:
|
|
warnings.simplefilter("always")
|
|
|
|
result = lx.extract(
|
|
text_or_documents="Sample text",
|
|
prompt_description="Extract entities",
|
|
examples=self.examples,
|
|
model_id="gemini-pro",
|
|
api_key="test_key",
|
|
use_schema_constraints=True,
|
|
)
|
|
|
|
# Check that a warning was issued
|
|
warning_messages = [str(warning.message) for warning in w]
|
|
self.assertTrue(
|
|
any("Gemini outputs native JSON" in msg for msg in warning_messages),
|
|
f"Expected Gemini-specific warning not found in: {warning_messages}",
|
|
)
|
|
|
|
# Result should still be returned
|
|
self.assertIsNotNone(result)
|
|
|
|
@mock.patch("langextract.factory.create_model")
|
|
def test_no_validation_without_schema(self, mock_create_model):
|
|
"""Test that validation is skipped when no schema is present."""
|
|
|
|
mock_model = mock.MagicMock()
|
|
mock_model.requires_fence_output = False
|
|
mock_model.schema = None # No schema
|
|
mock_model.infer.return_value = [
|
|
[mock.MagicMock(output='{"extractions": []}', score=1.0)]
|
|
]
|
|
|
|
mock_create_model.return_value = mock_model
|
|
|
|
with warnings.catch_warnings(record=True) as w:
|
|
warnings.simplefilter("always")
|
|
|
|
result = lx.extract(
|
|
text_or_documents="Sample text",
|
|
prompt_description="Extract",
|
|
examples=self.examples,
|
|
model_id="some-model",
|
|
api_key="key",
|
|
use_schema_constraints=False, # No schema constraints
|
|
)
|
|
|
|
# No format compatibility warnings should be issued
|
|
warning_messages = [str(warning.message) for warning in w]
|
|
self.assertFalse(
|
|
any("Format compatibility" in msg for msg in warning_messages),
|
|
f"Unexpected format warning found in: {warning_messages}",
|
|
)
|
|
|
|
self.assertIsNotNone(result)
|
|
|
|
|
|
@mock.patch.dict("os.environ", {"GEMINI_API_KEY": "test_key"})
|
|
class ExtractOutputSchemaTest(absltest.TestCase):
|
|
"""Tests for extract() with user-provided output_schema."""
|
|
|
|
def setUp(self):
|
|
super().setUp()
|
|
self.output_schema = lx.schema.extractions_schema(
|
|
lx.schema.extraction_item_schema("condition")
|
|
)
|
|
self.test_text = "Patient has hypertension"
|
|
|
|
def _patch_infer(self):
|
|
return mock.patch.object(
|
|
gemini.GeminiLanguageModel,
|
|
"infer",
|
|
autospec=True,
|
|
return_value=iter(
|
|
[[mock.Mock(output='{"extractions": [{"condition": "fever"}]}')]]
|
|
),
|
|
)
|
|
|
|
def test_extract_with_output_schema_allows_no_examples(self):
|
|
with self._patch_infer() as mock_infer:
|
|
result = lx.extract(
|
|
text_or_documents=self.test_text,
|
|
prompt_description="Extract conditions",
|
|
model_id="gemini-3.5-flash",
|
|
output_schema=self.output_schema,
|
|
)
|
|
|
|
model = mock_infer.call_args[0][0]
|
|
self.assertEqual(
|
|
model.schema.to_provider_config()["response_json_schema"],
|
|
self.output_schema,
|
|
)
|
|
self.assertLen(result.extractions, 1)
|
|
self.assertEqual(result.extractions[0].extraction_class, "condition")
|
|
self.assertEqual(result.extractions[0].extraction_text, "fever")
|
|
|
|
def test_extract_output_schema_overrides_example_schema(self):
|
|
examples = [
|
|
data.ExampleData(
|
|
text="Patient has diabetes",
|
|
extractions=[
|
|
data.Extraction(
|
|
extraction_class="condition",
|
|
extraction_text="diabetes",
|
|
)
|
|
],
|
|
)
|
|
]
|
|
|
|
with self._patch_infer() as mock_infer:
|
|
lx.extract(
|
|
text_or_documents=self.test_text,
|
|
prompt_description="Extract conditions",
|
|
examples=examples,
|
|
model_id="gemini-3.5-flash",
|
|
use_schema_constraints=True,
|
|
output_schema=self.output_schema,
|
|
)
|
|
|
|
model = mock_infer.call_args[0][0]
|
|
provider_config = model.schema.to_provider_config()
|
|
self.assertEqual(
|
|
provider_config["response_json_schema"], self.output_schema
|
|
)
|
|
self.assertNotIn("response_schema", provider_config)
|
|
|
|
def test_extract_requires_examples_without_output_schema(self):
|
|
with self.assertRaisesRegex(ValueError, "output_schema"):
|
|
lx.extract(
|
|
text_or_documents=self.test_text,
|
|
prompt_description="Extract conditions",
|
|
model_id="gemini-3.5-flash",
|
|
)
|
|
|
|
def test_extract_with_preconfigured_output_schema_model(self):
|
|
model = factory.create_model_from_id(
|
|
"gemini-3.5-flash", output_schema=self.output_schema
|
|
)
|
|
|
|
with self._patch_infer():
|
|
result = lx.extract(
|
|
text_or_documents=self.test_text,
|
|
prompt_description="Extract conditions",
|
|
model=model,
|
|
)
|
|
|
|
self.assertLen(result.extractions, 1)
|
|
|
|
def test_extract_applies_output_schema_to_plain_model(self):
|
|
model = factory.create_model_from_id("gemini-3.5-flash")
|
|
|
|
with self._patch_infer():
|
|
lx.extract(
|
|
text_or_documents=self.test_text,
|
|
prompt_description="Extract conditions",
|
|
model=model,
|
|
output_schema=self.output_schema,
|
|
)
|
|
|
|
self.assertTrue(model.schema.from_output_schema)
|
|
|
|
def test_extract_reapplies_same_output_schema_idempotently(self):
|
|
model = factory.create_model_from_id(
|
|
"gemini-3.5-flash", output_schema=self.output_schema
|
|
)
|
|
|
|
with self._patch_infer():
|
|
result = lx.extract(
|
|
text_or_documents=self.test_text,
|
|
prompt_description="Extract conditions",
|
|
model=model,
|
|
output_schema=self.output_schema,
|
|
)
|
|
|
|
self.assertIsNotNone(result)
|
|
|
|
def test_extract_output_schema_conflicts_with_example_schema_model(self):
|
|
examples = [
|
|
data.ExampleData(
|
|
text="Patient has diabetes",
|
|
extractions=[
|
|
data.Extraction(
|
|
extraction_class="condition",
|
|
extraction_text="diabetes",
|
|
)
|
|
],
|
|
)
|
|
]
|
|
model = factory.create_model(
|
|
factory.ModelConfig(model_id="gemini-3.5-flash"),
|
|
examples=examples,
|
|
use_schema_constraints=True,
|
|
)
|
|
|
|
with self.assertRaisesRegex(
|
|
exceptions.InferenceConfigError, "already has a schema"
|
|
):
|
|
lx.extract(
|
|
text_or_documents=self.test_text,
|
|
prompt_description="Extract conditions",
|
|
model=model,
|
|
output_schema=self.output_schema,
|
|
)
|
|
|
|
def test_extract_output_schema_rejects_fence_output(self):
|
|
with self.assertRaisesRegex(
|
|
exceptions.InferenceConfigError, "fence_output"
|
|
):
|
|
lx.extract(
|
|
text_or_documents=self.test_text,
|
|
prompt_description="Extract conditions",
|
|
model_id="gemini-3.5-flash",
|
|
output_schema=self.output_schema,
|
|
fence_output=True,
|
|
)
|
|
|
|
def test_extract_output_schema_rejects_yaml_format(self):
|
|
with self.assertRaisesRegex(
|
|
exceptions.InferenceConfigError, "format_type=JSON"
|
|
):
|
|
lx.extract(
|
|
text_or_documents=self.test_text,
|
|
prompt_description="Extract conditions",
|
|
model_id="gemini-3.5-flash",
|
|
output_schema=self.output_schema,
|
|
format_type=data.FormatType.YAML,
|
|
)
|
|
|
|
def test_extract_output_schema_rejects_fenced_resolver_params(self):
|
|
with self.assertRaisesRegex(
|
|
exceptions.InferenceConfigError, "fence_output"
|
|
):
|
|
lx.extract(
|
|
text_or_documents=self.test_text,
|
|
prompt_description="Extract conditions",
|
|
model_id="gemini-3.5-flash",
|
|
output_schema=self.output_schema,
|
|
resolver_params={"fence_output": True},
|
|
)
|
|
|
|
def test_extract_output_schema_rejects_unwrapped_resolver_output(self):
|
|
with self.assertRaisesRegex(exceptions.InferenceConfigError, "envelope"):
|
|
lx.extract(
|
|
text_or_documents=self.test_text,
|
|
prompt_description="Extract conditions",
|
|
model_id="gemini-3.5-flash",
|
|
output_schema=self.output_schema,
|
|
resolver_params={"require_extractions_key": False},
|
|
)
|
|
|
|
def test_extract_output_schema_rejects_custom_attribute_suffix(self):
|
|
with self.assertRaisesRegex(exceptions.InferenceConfigError, "envelope"):
|
|
lx.extract(
|
|
text_or_documents=self.test_text,
|
|
prompt_description="Extract conditions",
|
|
model_id="gemini-3.5-flash",
|
|
output_schema=self.output_schema,
|
|
resolver_params={"attribute_suffix": "_props"},
|
|
)
|
|
|
|
def test_extract_fence_conflict_leaves_caller_model_unmodified(self):
|
|
model = factory.create_model_from_id("gemini-3.5-flash")
|
|
|
|
with self.assertRaisesRegex(
|
|
exceptions.InferenceConfigError, "fence_output"
|
|
):
|
|
lx.extract(
|
|
text_or_documents=self.test_text,
|
|
prompt_description="Extract conditions",
|
|
model=model,
|
|
output_schema=self.output_schema,
|
|
fence_output=True,
|
|
)
|
|
|
|
self.assertIsNone(model.schema)
|
|
with self._patch_infer():
|
|
result = lx.extract(
|
|
text_or_documents=self.test_text,
|
|
prompt_description="Extract conditions",
|
|
model=model,
|
|
output_schema=self.output_schema,
|
|
)
|
|
self.assertIsNotNone(result)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
absltest.main()
|