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
284 lines
11 KiB
Python
284 lines
11 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.
|
|
|
|
import json
|
|
|
|
from absl.testing import absltest
|
|
from absl.testing import parameterized
|
|
import numpy as np
|
|
|
|
from langextract import data_lib
|
|
from langextract import io
|
|
from langextract.core import data
|
|
from langextract.core import tokenizer
|
|
|
|
|
|
class DataLibToDictParameterizedTest(parameterized.TestCase):
|
|
"""Tests conversion of AnnotatedDocument objects to JSON dicts.
|
|
|
|
Verifies that `annotated_document_to_dict` correctly serializes documents by:
|
|
- Excluding private fields (e.g., token_interval).
|
|
- Converting all expected extraction attributes properly.
|
|
- Handling int64 values for extraction indexes.
|
|
"""
|
|
|
|
@parameterized.named_parameters(
|
|
dict(
|
|
testcase_name="single_extraction_no_token_interval",
|
|
annotated_doc=data.AnnotatedDocument(
|
|
document_id="docA",
|
|
text="Just a short sentence.",
|
|
extractions=[
|
|
data.Extraction(
|
|
extraction_class="note",
|
|
extraction_text="short sentence",
|
|
extraction_index=1,
|
|
group_index=0,
|
|
),
|
|
],
|
|
),
|
|
expected_dict={
|
|
"document_id": "docA",
|
|
"extractions": [
|
|
{
|
|
"extraction_class": "note",
|
|
"extraction_text": "short sentence",
|
|
"char_interval": None,
|
|
"alignment_status": None,
|
|
"extraction_index": 1,
|
|
"group_index": 0,
|
|
"description": None,
|
|
"attributes": None,
|
|
},
|
|
],
|
|
"text": "Just a short sentence.",
|
|
},
|
|
),
|
|
dict(
|
|
testcase_name="multiple_extractions_with_token_interval",
|
|
annotated_doc=data.AnnotatedDocument(
|
|
document_id="docB",
|
|
text="Patient Jane reported a headache.",
|
|
extractions=[
|
|
data.Extraction(
|
|
extraction_class="patient",
|
|
extraction_text="Jane",
|
|
extraction_index=1,
|
|
group_index=0,
|
|
),
|
|
data.Extraction(
|
|
extraction_class="symptom",
|
|
extraction_text="headache",
|
|
extraction_index=2,
|
|
group_index=0,
|
|
char_interval=data.CharInterval(start_pos=24, end_pos=32),
|
|
token_interval=tokenizer.TokenInterval(
|
|
start_index=4, end_index=5
|
|
), # should be ignored
|
|
alignment_status=data.AlignmentStatus.MATCH_EXACT,
|
|
),
|
|
],
|
|
),
|
|
expected_dict={
|
|
"document_id": "docB",
|
|
"extractions": [
|
|
{
|
|
"extraction_class": "patient",
|
|
"extraction_text": "Jane",
|
|
"char_interval": None,
|
|
"alignment_status": None,
|
|
"extraction_index": 1,
|
|
"group_index": 0,
|
|
"description": None,
|
|
"attributes": None,
|
|
},
|
|
{
|
|
"extraction_class": "symptom",
|
|
"extraction_text": "headache",
|
|
"char_interval": {"start_pos": 24, "end_pos": 32},
|
|
"alignment_status": "match_exact",
|
|
"extraction_index": 2,
|
|
"group_index": 0,
|
|
"description": None,
|
|
"attributes": None,
|
|
},
|
|
],
|
|
"text": "Patient Jane reported a headache.",
|
|
},
|
|
),
|
|
dict(
|
|
testcase_name="extraction_with_attributes_and_token_interval",
|
|
annotated_doc=data.AnnotatedDocument(
|
|
document_id="docC",
|
|
text="He has mild chest pain and a cough.",
|
|
extractions=[
|
|
data.Extraction(
|
|
extraction_class="condition",
|
|
extraction_text="chest pain",
|
|
extraction_index=2,
|
|
group_index=1,
|
|
attributes={
|
|
"severity": "mild",
|
|
"persistence": "persistent",
|
|
},
|
|
char_interval=data.CharInterval(start_pos=12, end_pos=22),
|
|
token_interval=tokenizer.TokenInterval(
|
|
start_index=3, end_index=5
|
|
), # should be ignored
|
|
alignment_status=data.AlignmentStatus.MATCH_EXACT,
|
|
),
|
|
data.Extraction(
|
|
extraction_class="symptom",
|
|
extraction_text="cough",
|
|
extraction_index=3,
|
|
group_index=1,
|
|
),
|
|
],
|
|
),
|
|
expected_dict={
|
|
"document_id": "docC",
|
|
"extractions": [
|
|
{
|
|
"extraction_class": "condition",
|
|
"extraction_text": "chest pain",
|
|
"char_interval": {"start_pos": 12, "end_pos": 22},
|
|
"alignment_status": "match_exact",
|
|
"extraction_index": 2,
|
|
"group_index": 1,
|
|
"description": None,
|
|
"attributes": {
|
|
"severity": "mild",
|
|
"persistence": "persistent",
|
|
},
|
|
},
|
|
{
|
|
"extraction_class": "symptom",
|
|
"extraction_text": "cough",
|
|
"char_interval": None,
|
|
"alignment_status": None,
|
|
"extraction_index": 3,
|
|
"group_index": 1,
|
|
"description": None,
|
|
"attributes": None,
|
|
},
|
|
],
|
|
"text": "He has mild chest pain and a cough.",
|
|
},
|
|
),
|
|
)
|
|
def test_annotated_document_to_dict(self, annotated_doc, expected_dict):
|
|
actual_dict = data_lib.annotated_document_to_dict(annotated_doc)
|
|
self.assertDictEqual(
|
|
actual_dict,
|
|
expected_dict,
|
|
"annotated_document_to_dict() output differs from expected JSON dict.",
|
|
)
|
|
|
|
def test_annotated_document_to_dict_with_int64(self):
|
|
doc = data.AnnotatedDocument(
|
|
document_id="doc_int64",
|
|
text="Sample text with int64 index",
|
|
extractions=[
|
|
data.Extraction(
|
|
extraction_class="demo_extraction",
|
|
extraction_text="placeholder",
|
|
extraction_index=np.int64(42), # pytype: disable=wrong-arg-types
|
|
),
|
|
],
|
|
)
|
|
|
|
doc_dict = data_lib.annotated_document_to_dict(doc)
|
|
|
|
json_str = json.dumps(doc_dict, ensure_ascii=False)
|
|
self.assertIn('"extraction_index": 42', json_str)
|
|
|
|
|
|
class IsUrlTest(absltest.TestCase):
|
|
"""Tests for io.is_url function validation."""
|
|
|
|
def test_valid_urls(self):
|
|
"""Test that valid URLs are recognized."""
|
|
self.assertTrue(io.is_url("http://example.com"))
|
|
self.assertTrue(io.is_url("https://www.example.com"))
|
|
self.assertTrue(io.is_url("http://localhost:8080"))
|
|
self.assertTrue(io.is_url("http://192.168.1.1"))
|
|
self.assertTrue(io.is_url("http://[2001:db8::1]")) # IPv6
|
|
self.assertTrue(io.is_url("http://[::1]:8080")) # IPv6 localhost with port
|
|
|
|
def test_invalid_urls_with_text(self):
|
|
"""Test that URLs with additional text are rejected."""
|
|
# Validates fix for issue where text starting with URL was incorrectly fetched
|
|
self.assertFalse(io.is_url("http://example.com is a website"))
|
|
self.assertFalse(io.is_url("http://medical-journal.com published a study"))
|
|
|
|
def test_invalid_urls_no_scheme(self):
|
|
"""Test that URLs without proper scheme are rejected."""
|
|
self.assertFalse(io.is_url("example.com"))
|
|
self.assertFalse(io.is_url("www.example.com"))
|
|
self.assertFalse(io.is_url("ftp://example.com"))
|
|
|
|
|
|
class DocumentWithAdditionalContextTest(absltest.TestCase):
|
|
"""Tests for Document.with_additional_context()."""
|
|
|
|
def test_returns_new_document_with_overridden_context(self):
|
|
doc = data.Document(text="hello")
|
|
new_doc = doc.with_additional_context("ctx")
|
|
self.assertIsNot(new_doc, doc)
|
|
self.assertEqual(new_doc.additional_context, "ctx")
|
|
self.assertEqual(new_doc.text, "hello")
|
|
|
|
def test_does_not_mutate_original_context_or_tokenization(self):
|
|
doc = data.Document(text="hello")
|
|
doc.with_additional_context("ctx")
|
|
self.assertIsNone(doc.additional_context)
|
|
self.assertIsNone(doc._tokenized_text)
|
|
|
|
def test_preserves_explicit_document_id(self):
|
|
doc = data.Document(text="hello", document_id="custom-id")
|
|
new_doc = doc.with_additional_context("ctx")
|
|
self.assertEqual(new_doc.document_id, "custom-id")
|
|
|
|
def test_generates_shared_document_id_when_unset(self):
|
|
doc = data.Document(text="hello")
|
|
self.assertIsNone(doc._document_id)
|
|
|
|
new_doc = doc.with_additional_context("ctx")
|
|
|
|
self.assertIsNotNone(doc._document_id)
|
|
self.assertEqual(new_doc.document_id, doc.document_id)
|
|
|
|
def test_preserves_cached_tokenization(self):
|
|
doc = data.Document(text="hello world")
|
|
cached = doc.tokenized_text
|
|
new_doc = doc.with_additional_context("ctx")
|
|
self.assertIs(new_doc.tokenized_text, cached)
|
|
|
|
def test_preserves_both_explicit_id_and_cached_tokenization(self):
|
|
doc = data.Document(text="hello world", document_id="custom-id")
|
|
cached = doc.tokenized_text
|
|
new_doc = doc.with_additional_context("ctx")
|
|
self.assertEqual(new_doc.document_id, "custom-id")
|
|
self.assertIs(new_doc.tokenized_text, cached)
|
|
|
|
def test_accepts_none_context(self):
|
|
doc = data.Document(text="hello", additional_context="original")
|
|
new_doc = doc.with_additional_context(None)
|
|
self.assertIsNone(new_doc.additional_context)
|
|
self.assertEqual(new_doc.text, "hello")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
absltest.main()
|