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
596 lines
20 KiB
Python
596 lines
20 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 textwrap
|
|
|
|
from absl.testing import absltest
|
|
from absl.testing import parameterized
|
|
|
|
from langextract import prompting
|
|
from langextract.core import data
|
|
from langextract.core import format_handler as fh
|
|
|
|
|
|
class QAPromptGeneratorTest(parameterized.TestCase):
|
|
|
|
def test_generate_prompt(self):
|
|
prompt_template_structured = prompting.PromptTemplateStructured(
|
|
description=(
|
|
"You are an assistant specialized in extracting key extractions"
|
|
" from text.\nIdentify and extract important extractions such as"
|
|
" people, places,\norganizations, dates, and medical conditions"
|
|
" mentioned in the text.\n**Please ensure that the extractions are"
|
|
" extracted in the same order as they\nappear in the source"
|
|
" text.**\nProvide the extracted extractions in a structured YAML"
|
|
" format."
|
|
),
|
|
examples=[
|
|
data.ExampleData(
|
|
text=(
|
|
"The patient was diagnosed with hypertension and diabetes."
|
|
),
|
|
extractions=[
|
|
data.Extraction(
|
|
extraction_text="hypertension",
|
|
extraction_class="medical_condition",
|
|
attributes={
|
|
"chronicity": "chronic",
|
|
"system": "cardiovascular",
|
|
},
|
|
),
|
|
data.Extraction(
|
|
extraction_text="diabetes",
|
|
extraction_class="medical_condition",
|
|
attributes={
|
|
"chronicity": "chronic",
|
|
"system": "endocrine",
|
|
},
|
|
),
|
|
],
|
|
)
|
|
],
|
|
)
|
|
|
|
format_handler = fh.FormatHandler(
|
|
format_type=data.FormatType.YAML,
|
|
use_wrapper=True,
|
|
wrapper_key="extractions",
|
|
use_fences=True,
|
|
)
|
|
|
|
prompt_generator = prompting.QAPromptGenerator(
|
|
template=prompt_template_structured,
|
|
format_handler=format_handler,
|
|
examples_heading="",
|
|
question_prefix="",
|
|
answer_prefix="",
|
|
)
|
|
|
|
actual_prompt_text = prompt_generator.render(
|
|
"The patient reports chest pain and shortness of breath."
|
|
)
|
|
|
|
expected_prompt_text = textwrap.dedent(f"""\
|
|
You are an assistant specialized in extracting key extractions from text.
|
|
Identify and extract important extractions such as people, places,
|
|
organizations, dates, and medical conditions mentioned in the text.
|
|
**Please ensure that the extractions are extracted in the same order as they
|
|
appear in the source text.**
|
|
Provide the extracted extractions in a structured YAML format.
|
|
|
|
|
|
The patient was diagnosed with hypertension and diabetes.
|
|
```yaml
|
|
{data.EXTRACTIONS_KEY}:
|
|
- medical_condition: hypertension
|
|
medical_condition_attributes:
|
|
chronicity: chronic
|
|
system: cardiovascular
|
|
- medical_condition: diabetes
|
|
medical_condition_attributes:
|
|
chronicity: chronic
|
|
system: endocrine
|
|
```
|
|
|
|
The patient reports chest pain and shortness of breath.
|
|
""")
|
|
self.assertEqual(expected_prompt_text, actual_prompt_text)
|
|
|
|
@parameterized.named_parameters(
|
|
dict(
|
|
testcase_name="json_basic_format",
|
|
format_type=data.FormatType.JSON,
|
|
example_text="Patient has diabetes and is prescribed insulin.",
|
|
example_extractions=[
|
|
data.Extraction(
|
|
extraction_text="diabetes",
|
|
extraction_class="medical_condition",
|
|
attributes={"chronicity": "chronic"},
|
|
),
|
|
data.Extraction(
|
|
extraction_text="insulin",
|
|
extraction_class="medication",
|
|
attributes={"prescribed": "prescribed"},
|
|
),
|
|
],
|
|
expected_formatted_example=textwrap.dedent(f"""\
|
|
Patient has diabetes and is prescribed insulin.
|
|
```json
|
|
{{
|
|
"{data.EXTRACTIONS_KEY}": [
|
|
{{
|
|
"medical_condition": "diabetes",
|
|
"medical_condition_attributes": {{
|
|
"chronicity": "chronic"
|
|
}}
|
|
}},
|
|
{{
|
|
"medication": "insulin",
|
|
"medication_attributes": {{
|
|
"prescribed": "prescribed"
|
|
}}
|
|
}}
|
|
]
|
|
}}
|
|
```
|
|
"""),
|
|
),
|
|
dict(
|
|
testcase_name="yaml_basic_format",
|
|
format_type=data.FormatType.YAML,
|
|
example_text="Patient has diabetes and is prescribed insulin.",
|
|
example_extractions=[
|
|
data.Extraction(
|
|
extraction_text="diabetes",
|
|
extraction_class="medical_condition",
|
|
attributes={"chronicity": "chronic"},
|
|
),
|
|
data.Extraction(
|
|
extraction_text="insulin",
|
|
extraction_class="medication",
|
|
attributes={"prescribed": "prescribed"},
|
|
),
|
|
],
|
|
expected_formatted_example=textwrap.dedent(f"""\
|
|
Patient has diabetes and is prescribed insulin.
|
|
```yaml
|
|
{data.EXTRACTIONS_KEY}:
|
|
- medical_condition: diabetes
|
|
medical_condition_attributes:
|
|
chronicity: chronic
|
|
- medication: insulin
|
|
medication_attributes:
|
|
prescribed: prescribed
|
|
```
|
|
"""),
|
|
),
|
|
dict(
|
|
testcase_name="custom_attribute_suffix",
|
|
format_type=data.FormatType.YAML,
|
|
example_text="Patient has a fever.",
|
|
example_extractions=[
|
|
data.Extraction(
|
|
extraction_text="fever",
|
|
extraction_class="symptom",
|
|
attributes={"severity": "mild"},
|
|
),
|
|
],
|
|
attribute_suffix="_props",
|
|
expected_formatted_example=textwrap.dedent(f"""\
|
|
Patient has a fever.
|
|
```yaml
|
|
{data.EXTRACTIONS_KEY}:
|
|
- symptom: fever
|
|
symptom_props:
|
|
severity: mild
|
|
```
|
|
"""),
|
|
),
|
|
dict(
|
|
testcase_name="yaml_empty_extractions",
|
|
format_type=data.FormatType.YAML,
|
|
example_text="Text with no extractions.",
|
|
example_extractions=[],
|
|
expected_formatted_example=textwrap.dedent(f"""\
|
|
Text with no extractions.
|
|
```yaml
|
|
{data.EXTRACTIONS_KEY}: []
|
|
```
|
|
"""),
|
|
),
|
|
dict(
|
|
testcase_name="json_empty_extractions",
|
|
format_type=data.FormatType.JSON,
|
|
example_text="Text with no extractions.",
|
|
example_extractions=[],
|
|
expected_formatted_example=textwrap.dedent(f"""\
|
|
Text with no extractions.
|
|
```json
|
|
{{
|
|
"{data.EXTRACTIONS_KEY}": []
|
|
}}
|
|
```
|
|
"""),
|
|
),
|
|
dict(
|
|
testcase_name="yaml_empty_attributes",
|
|
format_type=data.FormatType.YAML,
|
|
example_text="Patient is resting comfortably.",
|
|
example_extractions=[
|
|
data.Extraction(
|
|
extraction_text="Patient",
|
|
extraction_class="person",
|
|
attributes={},
|
|
),
|
|
],
|
|
expected_formatted_example=textwrap.dedent(f"""\
|
|
Patient is resting comfortably.
|
|
```yaml
|
|
{data.EXTRACTIONS_KEY}:
|
|
- person: Patient
|
|
person_attributes: {{}}
|
|
```
|
|
"""),
|
|
),
|
|
dict(
|
|
testcase_name="json_empty_attributes",
|
|
format_type=data.FormatType.JSON,
|
|
example_text="Patient is resting comfortably.",
|
|
example_extractions=[
|
|
data.Extraction(
|
|
extraction_text="Patient",
|
|
extraction_class="person",
|
|
attributes={},
|
|
),
|
|
],
|
|
expected_formatted_example=textwrap.dedent(f"""\
|
|
Patient is resting comfortably.
|
|
```json
|
|
{{
|
|
"{data.EXTRACTIONS_KEY}": [
|
|
{{
|
|
"person": "Patient",
|
|
"person_attributes": {{}}
|
|
}}
|
|
]
|
|
}}
|
|
```
|
|
"""),
|
|
),
|
|
dict(
|
|
testcase_name="yaml_same_extraction_class_multiple_times",
|
|
format_type=data.FormatType.YAML,
|
|
example_text=(
|
|
"Patient has multiple medications: aspirin and lisinopril."
|
|
),
|
|
example_extractions=[
|
|
data.Extraction(
|
|
extraction_text="aspirin",
|
|
extraction_class="medication",
|
|
attributes={"dosage": "81mg"},
|
|
),
|
|
data.Extraction(
|
|
extraction_text="lisinopril",
|
|
extraction_class="medication",
|
|
attributes={"dosage": "10mg"},
|
|
),
|
|
],
|
|
expected_formatted_example=textwrap.dedent(f"""\
|
|
Patient has multiple medications: aspirin and lisinopril.
|
|
```yaml
|
|
{data.EXTRACTIONS_KEY}:
|
|
- medication: aspirin
|
|
medication_attributes:
|
|
dosage: 81mg
|
|
- medication: lisinopril
|
|
medication_attributes:
|
|
dosage: 10mg
|
|
```
|
|
"""),
|
|
),
|
|
dict(
|
|
testcase_name="json_simplified_no_extractions_key",
|
|
format_type=data.FormatType.JSON,
|
|
example_text="Patient has diabetes and is prescribed insulin.",
|
|
example_extractions=[
|
|
data.Extraction(
|
|
extraction_text="diabetes",
|
|
extraction_class="medical_condition",
|
|
attributes={"chronicity": "chronic"},
|
|
),
|
|
data.Extraction(
|
|
extraction_text="insulin",
|
|
extraction_class="medication",
|
|
attributes={"prescribed": "prescribed"},
|
|
),
|
|
],
|
|
require_extractions_key=False,
|
|
expected_formatted_example=textwrap.dedent("""\
|
|
Patient has diabetes and is prescribed insulin.
|
|
```json
|
|
[
|
|
{
|
|
"medical_condition": "diabetes",
|
|
"medical_condition_attributes": {
|
|
"chronicity": "chronic"
|
|
}
|
|
},
|
|
{
|
|
"medication": "insulin",
|
|
"medication_attributes": {
|
|
"prescribed": "prescribed"
|
|
}
|
|
}
|
|
]
|
|
```
|
|
"""),
|
|
),
|
|
dict(
|
|
testcase_name="yaml_simplified_no_extractions_key",
|
|
format_type=data.FormatType.YAML,
|
|
example_text="Patient has a fever.",
|
|
example_extractions=[
|
|
data.Extraction(
|
|
extraction_text="fever",
|
|
extraction_class="symptom",
|
|
attributes={"severity": "mild"},
|
|
),
|
|
],
|
|
require_extractions_key=False,
|
|
expected_formatted_example=textwrap.dedent("""\
|
|
Patient has a fever.
|
|
```yaml
|
|
- symptom: fever
|
|
symptom_attributes:
|
|
severity: mild
|
|
```
|
|
"""),
|
|
),
|
|
)
|
|
def test_format_example(
|
|
self,
|
|
format_type,
|
|
example_text,
|
|
example_extractions,
|
|
expected_formatted_example,
|
|
attribute_suffix="_attributes",
|
|
require_extractions_key=True,
|
|
):
|
|
"""Tests formatting of examples in different formats and scenarios."""
|
|
example_data = data.ExampleData(
|
|
text=example_text,
|
|
extractions=example_extractions,
|
|
)
|
|
|
|
structured_template = prompting.PromptTemplateStructured(
|
|
description="Extract information from the text.",
|
|
examples=[example_data],
|
|
)
|
|
|
|
format_handler = fh.FormatHandler(
|
|
format_type=format_type,
|
|
use_wrapper=require_extractions_key,
|
|
wrapper_key="extractions" if require_extractions_key else None,
|
|
use_fences=True,
|
|
attribute_suffix=attribute_suffix,
|
|
)
|
|
|
|
prompt_generator = prompting.QAPromptGenerator(
|
|
template=structured_template,
|
|
format_handler=format_handler,
|
|
question_prefix="",
|
|
answer_prefix="",
|
|
)
|
|
|
|
actual_formatted_example = prompt_generator.format_example_as_text(
|
|
example_data
|
|
)
|
|
self.assertEqual(expected_formatted_example, actual_formatted_example)
|
|
|
|
|
|
class PromptBuilderTest(absltest.TestCase):
|
|
"""Tests for PromptBuilder base class."""
|
|
|
|
def _create_generator(self):
|
|
"""Creates a simple QAPromptGenerator for testing."""
|
|
template = prompting.PromptTemplateStructured(
|
|
description="Extract entities.",
|
|
examples=[
|
|
data.ExampleData(
|
|
text="Sample text.",
|
|
extractions=[
|
|
data.Extraction(
|
|
extraction_text="Sample",
|
|
extraction_class="entity",
|
|
)
|
|
],
|
|
)
|
|
],
|
|
)
|
|
format_handler = fh.FormatHandler(
|
|
format_type=data.FormatType.YAML,
|
|
use_wrapper=True,
|
|
wrapper_key="extractions",
|
|
use_fences=True,
|
|
)
|
|
return prompting.QAPromptGenerator(
|
|
template=template,
|
|
format_handler=format_handler,
|
|
)
|
|
|
|
def test_build_prompt_renders_chunk_text(self):
|
|
"""Verifies build_prompt includes chunk text in the rendered prompt."""
|
|
generator = self._create_generator()
|
|
builder = prompting.PromptBuilder(generator)
|
|
|
|
prompt = builder.build_prompt(
|
|
chunk_text="Test input text.",
|
|
document_id="doc1",
|
|
)
|
|
|
|
self.assertIn("Test input text.", prompt)
|
|
self.assertIn("Extract entities.", prompt)
|
|
|
|
def test_build_prompt_includes_additional_context(self):
|
|
"""Verifies build_prompt passes additional_context to renderer."""
|
|
generator = self._create_generator()
|
|
builder = prompting.PromptBuilder(generator)
|
|
|
|
prompt = builder.build_prompt(
|
|
chunk_text="Test input.",
|
|
document_id="doc1",
|
|
additional_context="Important context here.",
|
|
)
|
|
|
|
self.assertIn("Important context here.", prompt)
|
|
|
|
|
|
class ContextAwarePromptBuilderTest(absltest.TestCase):
|
|
"""Tests for ContextAwarePromptBuilder."""
|
|
|
|
def _create_generator(self):
|
|
"""Creates a simple QAPromptGenerator for testing."""
|
|
template = prompting.PromptTemplateStructured(
|
|
description="Extract entities.",
|
|
examples=[
|
|
data.ExampleData(
|
|
text="Sample text.",
|
|
extractions=[
|
|
data.Extraction(
|
|
extraction_text="Sample",
|
|
extraction_class="entity",
|
|
)
|
|
],
|
|
)
|
|
],
|
|
)
|
|
format_handler = fh.FormatHandler(
|
|
format_type=data.FormatType.YAML,
|
|
use_wrapper=True,
|
|
wrapper_key="extractions",
|
|
use_fences=True,
|
|
)
|
|
return prompting.QAPromptGenerator(
|
|
template=template,
|
|
format_handler=format_handler,
|
|
)
|
|
|
|
def test_context_window_chars_property(self):
|
|
"""Verifies the context_window_chars property returns configured value."""
|
|
generator = self._create_generator()
|
|
|
|
builder_none = prompting.ContextAwarePromptBuilder(generator)
|
|
self.assertIsNone(builder_none.context_window_chars)
|
|
|
|
builder_with_value = prompting.ContextAwarePromptBuilder(
|
|
generator, context_window_chars=100
|
|
)
|
|
self.assertEqual(100, builder_with_value.context_window_chars)
|
|
|
|
def test_first_chunk_has_no_previous_context(self):
|
|
"""Verifies the first chunk does not include previous context."""
|
|
generator = self._create_generator()
|
|
builder = prompting.ContextAwarePromptBuilder(
|
|
generator, context_window_chars=50
|
|
)
|
|
context_prefix = prompting.ContextAwarePromptBuilder._CONTEXT_PREFIX
|
|
|
|
prompt = builder.build_prompt(
|
|
chunk_text="First chunk text.",
|
|
document_id="doc1",
|
|
)
|
|
|
|
self.assertNotIn(context_prefix, prompt)
|
|
self.assertIn("First chunk text.", prompt)
|
|
|
|
def test_second_chunk_includes_previous_context(self):
|
|
"""Verifies the second chunk includes text from the first chunk."""
|
|
generator = self._create_generator()
|
|
builder = prompting.ContextAwarePromptBuilder(
|
|
generator, context_window_chars=20
|
|
)
|
|
context_prefix = prompting.ContextAwarePromptBuilder._CONTEXT_PREFIX
|
|
|
|
builder.build_prompt(chunk_text="First chunk ending.", document_id="doc1")
|
|
second_prompt = builder.build_prompt(
|
|
chunk_text="Second chunk text.",
|
|
document_id="doc1",
|
|
)
|
|
|
|
self.assertIn(context_prefix, second_prompt)
|
|
self.assertIn("chunk ending.", second_prompt)
|
|
|
|
def test_context_disabled_when_none(self):
|
|
"""Verifies no context is added when context_window_chars is None."""
|
|
generator = self._create_generator()
|
|
builder = prompting.ContextAwarePromptBuilder(
|
|
generator, context_window_chars=None
|
|
)
|
|
context_prefix = prompting.ContextAwarePromptBuilder._CONTEXT_PREFIX
|
|
|
|
builder.build_prompt(chunk_text="First chunk.", document_id="doc1")
|
|
second_prompt = builder.build_prompt(
|
|
chunk_text="Second chunk.",
|
|
document_id="doc1",
|
|
)
|
|
|
|
self.assertNotIn(context_prefix, second_prompt)
|
|
|
|
def test_context_isolated_per_document(self):
|
|
"""Verifies context tracking is isolated per document_id."""
|
|
generator = self._create_generator()
|
|
builder = prompting.ContextAwarePromptBuilder(
|
|
generator, context_window_chars=50
|
|
)
|
|
|
|
builder.build_prompt(chunk_text="Doc A chunk one.", document_id="docA")
|
|
builder.build_prompt(chunk_text="Doc B chunk one.", document_id="docB")
|
|
|
|
prompt_a2 = builder.build_prompt(
|
|
chunk_text="Doc A chunk two.",
|
|
document_id="docA",
|
|
)
|
|
prompt_b2 = builder.build_prompt(
|
|
chunk_text="Doc B chunk two.",
|
|
document_id="docB",
|
|
)
|
|
|
|
self.assertIn("Doc A chunk one", prompt_a2)
|
|
self.assertNotIn("Doc B", prompt_a2)
|
|
self.assertIn("Doc B chunk one", prompt_b2)
|
|
self.assertNotIn("Doc A", prompt_b2)
|
|
|
|
def test_combines_previous_context_with_additional_context(self):
|
|
"""Verifies both previous chunk context and additional_context are included."""
|
|
generator = self._create_generator()
|
|
builder = prompting.ContextAwarePromptBuilder(
|
|
generator, context_window_chars=30
|
|
)
|
|
context_prefix = prompting.ContextAwarePromptBuilder._CONTEXT_PREFIX
|
|
|
|
builder.build_prompt(chunk_text="Previous chunk text.", document_id="doc1")
|
|
prompt = builder.build_prompt(
|
|
chunk_text="Current chunk.",
|
|
document_id="doc1",
|
|
additional_context="Extra info here.",
|
|
)
|
|
|
|
self.assertIn(context_prefix, prompt)
|
|
self.assertIn("Previous chunk text.", prompt)
|
|
self.assertIn("Extra info here.", prompt)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
absltest.main()
|