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
1260 lines
40 KiB
Python
1260 lines
40 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.
|
|
|
|
"""Live API integration tests that require real API keys.
|
|
|
|
These tests are skipped if API keys are not available in the environment.
|
|
They should run in CI after all other tests pass.
|
|
"""
|
|
|
|
import functools
|
|
import json
|
|
import os
|
|
import re
|
|
import textwrap
|
|
import time
|
|
from typing import Any
|
|
import unittest
|
|
from unittest import mock
|
|
import uuid
|
|
|
|
import dotenv
|
|
import google.auth
|
|
import google.auth.exceptions
|
|
import google.genai.errors
|
|
import pytest
|
|
|
|
from langextract import data
|
|
import langextract as lx
|
|
from langextract.core import tokenizer as tokenizer_lib
|
|
from langextract.providers import gemini_batch as gb
|
|
from langextract.providers import openai_batch
|
|
|
|
dotenv.load_dotenv(override=True)
|
|
|
|
DEFAULT_GEMINI_MODEL = "gemini-3.5-flash"
|
|
DEFAULT_OPENAI_MODEL = "gpt-4o"
|
|
|
|
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY") or os.environ.get(
|
|
"LANGEXTRACT_API_KEY"
|
|
)
|
|
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
|
|
RUN_OPENAI_BATCH_LIVE_TESTS = (
|
|
os.environ.get("LANGEXTRACT_RUN_OPENAI_BATCH_LIVE_TESTS") == "1"
|
|
)
|
|
|
|
VERTEX_PROJECT = os.environ.get("VERTEX_PROJECT") or os.environ.get(
|
|
"GOOGLE_CLOUD_PROJECT"
|
|
)
|
|
VERTEX_LOCATION = os.environ.get("VERTEX_LOCATION", "us-central1")
|
|
|
|
|
|
def has_vertex_ai_credentials():
|
|
"""Check if Vertex AI credentials are available."""
|
|
if not VERTEX_PROJECT:
|
|
return False
|
|
try:
|
|
credentials, _ = google.auth.default()
|
|
return credentials is not None
|
|
except (ImportError, google.auth.exceptions.DefaultCredentialsError):
|
|
return False
|
|
|
|
|
|
skip_if_no_gemini = pytest.mark.skipif(
|
|
not GEMINI_API_KEY,
|
|
reason=(
|
|
"Gemini API key not available (set GEMINI_API_KEY or"
|
|
" LANGEXTRACT_API_KEY)"
|
|
),
|
|
)
|
|
skip_if_no_openai = pytest.mark.skipif(
|
|
not OPENAI_API_KEY,
|
|
reason="OpenAI API key not available (set OPENAI_API_KEY)",
|
|
)
|
|
skip_if_openai_batch_live_disabled = pytest.mark.skipif(
|
|
not RUN_OPENAI_BATCH_LIVE_TESTS,
|
|
reason=(
|
|
"OpenAI Batch API live test not enabled "
|
|
"(set LANGEXTRACT_RUN_OPENAI_BATCH_LIVE_TESTS=1)"
|
|
),
|
|
)
|
|
skip_if_no_vertex = pytest.mark.skipif(
|
|
not has_vertex_ai_credentials(),
|
|
reason=(
|
|
"Vertex AI credentials not available (set GOOGLE_CLOUD_PROJECT and"
|
|
" configure gcloud auth)"
|
|
),
|
|
)
|
|
|
|
live_api = pytest.mark.live_api
|
|
|
|
GEMINI_MODEL_PARAMS = {
|
|
"temperature": 0.0,
|
|
"top_p": 0.0,
|
|
"max_output_tokens": 256,
|
|
}
|
|
|
|
OPENAI_MODEL_PARAMS = {
|
|
"temperature": 0.0,
|
|
}
|
|
|
|
# Extraction Classes
|
|
_CLASS_MEDICATION = "medication"
|
|
_CLASS_DOSAGE = "dosage"
|
|
_CLASS_ROUTE = "route"
|
|
_CLASS_FREQUENCY = "frequency"
|
|
_CLASS_DURATION = "duration"
|
|
_CLASS_CONDITION = "condition"
|
|
|
|
INITIAL_RETRY_DELAY = 1.0
|
|
MAX_RETRY_DELAY = 8.0
|
|
|
|
|
|
def retry_on_transient_errors(max_retries=3, backoff_factor=2.0):
|
|
"""Decorator to retry tests on transient API errors with exponential backoff.
|
|
|
|
Args:
|
|
max_retries (int): Maximum number of retry attempts
|
|
backoff_factor (float): Multiplier for exponential backoff (e.g., 2.0 = 1s, 2s, 4s)
|
|
"""
|
|
|
|
def decorator(func):
|
|
@functools.wraps(func)
|
|
def wrapper(*args, **kwargs):
|
|
last_exception = None
|
|
delay = INITIAL_RETRY_DELAY
|
|
|
|
for attempt in range(max_retries + 1):
|
|
try:
|
|
return func(*args, **kwargs)
|
|
except (
|
|
lx.exceptions.LangExtractError,
|
|
google.genai.errors.ClientError,
|
|
ConnectionError,
|
|
TimeoutError,
|
|
OSError,
|
|
RuntimeError,
|
|
) as e:
|
|
last_exception = e
|
|
if attempt < max_retries:
|
|
print(
|
|
f"\nRetryable error ({type(e).__name__}) on attempt"
|
|
f" {attempt + 1}/{max_retries + 1}: {e}"
|
|
)
|
|
time.sleep(delay)
|
|
delay = min(delay * backoff_factor, MAX_RETRY_DELAY)
|
|
continue
|
|
|
|
raise
|
|
|
|
raise last_exception
|
|
|
|
return wrapper
|
|
|
|
return decorator
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def add_delay_between_tests():
|
|
"""Add a small delay between tests to avoid rate limiting."""
|
|
yield
|
|
time.sleep(0.5)
|
|
|
|
|
|
def get_basic_medication_examples():
|
|
"""Get example data for basic medication extraction."""
|
|
return [
|
|
lx.data.ExampleData(
|
|
text="Patient was given 250 mg IV Cefazolin TID for one week.",
|
|
extractions=[
|
|
lx.data.Extraction(
|
|
extraction_class=_CLASS_DOSAGE, extraction_text="250 mg"
|
|
),
|
|
lx.data.Extraction(
|
|
extraction_class=_CLASS_ROUTE, extraction_text="IV"
|
|
),
|
|
lx.data.Extraction(
|
|
extraction_class=_CLASS_MEDICATION,
|
|
extraction_text="Cefazolin",
|
|
),
|
|
lx.data.Extraction(
|
|
extraction_class=_CLASS_FREQUENCY,
|
|
extraction_text="TID", # TID = three times a day
|
|
),
|
|
lx.data.Extraction(
|
|
extraction_class=_CLASS_DURATION,
|
|
extraction_text="for one week",
|
|
),
|
|
],
|
|
)
|
|
]
|
|
|
|
|
|
def get_relationship_examples():
|
|
"""Get example data for medication relationship extraction."""
|
|
return [
|
|
lx.data.ExampleData(
|
|
text=(
|
|
"Patient takes Aspirin 100mg daily for heart health and"
|
|
" Simvastatin 20mg at bedtime."
|
|
),
|
|
extractions=[
|
|
# First medication group
|
|
lx.data.Extraction(
|
|
extraction_class=_CLASS_MEDICATION,
|
|
extraction_text="Aspirin",
|
|
attributes={"medication_group": "Aspirin"},
|
|
),
|
|
lx.data.Extraction(
|
|
extraction_class=_CLASS_DOSAGE,
|
|
extraction_text="100mg",
|
|
attributes={"medication_group": "Aspirin"},
|
|
),
|
|
lx.data.Extraction(
|
|
extraction_class=_CLASS_FREQUENCY,
|
|
extraction_text="daily",
|
|
attributes={"medication_group": "Aspirin"},
|
|
),
|
|
lx.data.Extraction(
|
|
extraction_class=_CLASS_CONDITION,
|
|
extraction_text="heart health",
|
|
attributes={"medication_group": "Aspirin"},
|
|
),
|
|
# Second medication group
|
|
lx.data.Extraction(
|
|
extraction_class=_CLASS_MEDICATION,
|
|
extraction_text="Simvastatin",
|
|
attributes={"medication_group": "Simvastatin"},
|
|
),
|
|
lx.data.Extraction(
|
|
extraction_class=_CLASS_DOSAGE,
|
|
extraction_text="20mg",
|
|
attributes={"medication_group": "Simvastatin"},
|
|
),
|
|
lx.data.Extraction(
|
|
extraction_class=_CLASS_FREQUENCY,
|
|
extraction_text="at bedtime",
|
|
attributes={"medication_group": "Simvastatin"},
|
|
),
|
|
],
|
|
)
|
|
]
|
|
|
|
|
|
def extract_by_class(result, extraction_class):
|
|
"""Helper to extract entities by class.
|
|
|
|
Returns a set of extraction texts for the given class.
|
|
"""
|
|
return {
|
|
e.extraction_text
|
|
for e in result.extractions
|
|
if e.extraction_class == extraction_class
|
|
}
|
|
|
|
|
|
def assert_extractions_contain(test_case, result, expected_classes):
|
|
"""Assert that result contains all expected extraction classes.
|
|
|
|
Uses unittest assertions for richer error messages.
|
|
"""
|
|
actual_classes = {e.extraction_class for e in result.extractions}
|
|
missing_classes = expected_classes - actual_classes
|
|
test_case.assertFalse(
|
|
missing_classes,
|
|
f"Missing expected classes: {missing_classes}. Found extractions:"
|
|
f" {[f'{e.extraction_class}:{e.extraction_text}' for e in result.extractions]}",
|
|
)
|
|
|
|
|
|
def assert_valid_char_intervals(test_case, result):
|
|
"""Assert that all extractions have valid char intervals and alignment status."""
|
|
for extraction in result.extractions:
|
|
test_case.assertIsNotNone(
|
|
extraction.char_interval,
|
|
f"Missing char_interval for extraction: {extraction.extraction_text}",
|
|
)
|
|
test_case.assertIsNotNone(
|
|
extraction.alignment_status,
|
|
"Missing alignment_status for extraction:"
|
|
f" {extraction.extraction_text}",
|
|
)
|
|
if isinstance(result, lx.data.AnnotatedDocument) and result.text:
|
|
text_length = len(result.text)
|
|
test_case.assertGreaterEqual(
|
|
extraction.char_interval.start_pos,
|
|
0,
|
|
f"Invalid start_pos for extraction: {extraction.extraction_text}",
|
|
)
|
|
test_case.assertLessEqual(
|
|
extraction.char_interval.end_pos,
|
|
text_length,
|
|
f"Invalid end_pos for extraction: {extraction.extraction_text}",
|
|
)
|
|
|
|
|
|
class TestLiveAPIGemini(unittest.TestCase):
|
|
"""Tests using real Gemini API."""
|
|
|
|
def _check_cached_result(self, result_json: dict[str, Any]) -> bool:
|
|
"""Check if cached result contains expected medication data.
|
|
|
|
Args:
|
|
result_json: The raw JSON dict from the cache file.
|
|
Expected format: {"text": "JSON_STRING_OF_RESULT"}
|
|
|
|
Returns:
|
|
True if the result contains valid medication extractions, False otherwise.
|
|
"""
|
|
try:
|
|
text_content = result_json.get("text")
|
|
if not isinstance(text_content, str):
|
|
return False
|
|
|
|
inner_json = json.loads(text_content)
|
|
if not isinstance(inner_json, dict):
|
|
return False
|
|
|
|
extractions_data = inner_json.get(data.EXTRACTIONS_KEY)
|
|
if not isinstance(extractions_data, list):
|
|
return False
|
|
|
|
extractions = []
|
|
for item in extractions_data:
|
|
if isinstance(item, dict):
|
|
clean_item = {k: v for k, v in item.items() if not k.startswith("_")}
|
|
extractions.append(data.Extraction(**clean_item))
|
|
|
|
doc = data.AnnotatedDocument(
|
|
text=inner_json.get("text"), extractions=extractions
|
|
)
|
|
|
|
if not doc.extractions:
|
|
return False
|
|
|
|
# Check for specific content
|
|
medication_texts = extract_by_class(doc, _CLASS_MEDICATION)
|
|
dosage_texts = extract_by_class(doc, _CLASS_DOSAGE)
|
|
|
|
has_lisinopril = any("Lisinopril" in t for t in medication_texts)
|
|
has_10mg = any("10mg" in t for t in dosage_texts)
|
|
|
|
return has_lisinopril and has_10mg
|
|
|
|
except (json.JSONDecodeError, TypeError, ValueError):
|
|
return False
|
|
|
|
def _verify_gcs_cache_content(self, bucket_name):
|
|
"""Verify that GCS cache contains expected structured results."""
|
|
cache = gb.GCSBatchCache(bucket_name, project=VERTEX_PROJECT)
|
|
found_content = False
|
|
|
|
# Use iter_items() to check cache content
|
|
items = list(cache.iter_items())
|
|
self.assertTrue(len(items) > 0, "No cache files found in GCS bucket")
|
|
|
|
for _, text in items:
|
|
try:
|
|
result_json = json.loads(text)
|
|
if self._check_cached_result(result_json):
|
|
found_content = True
|
|
break
|
|
except (json.JSONDecodeError, TypeError, ValueError):
|
|
continue
|
|
|
|
self.assertTrue(
|
|
found_content,
|
|
"Could not find expected structured result in GCS cache files",
|
|
)
|
|
|
|
@skip_if_no_gemini
|
|
@live_api
|
|
@retry_on_transient_errors(max_retries=2)
|
|
def test_medication_extraction(self):
|
|
"""Test medication extraction with entities in order."""
|
|
prompt = textwrap.dedent("""\
|
|
Extract medication information including medication name, dosage, route, frequency,
|
|
and duration in the order they appear in the text.""")
|
|
|
|
examples = get_basic_medication_examples()
|
|
input_text = "Patient took 400 mg PO Ibuprofen q4h for two days."
|
|
|
|
result = lx.extract(
|
|
text_or_documents=input_text,
|
|
prompt_description=prompt,
|
|
examples=examples,
|
|
model_id=DEFAULT_GEMINI_MODEL,
|
|
api_key=GEMINI_API_KEY,
|
|
language_model_params=GEMINI_MODEL_PARAMS,
|
|
)
|
|
|
|
assert result is not None
|
|
self.assertIsInstance(result, lx.data.AnnotatedDocument)
|
|
assert len(result.extractions) > 0
|
|
|
|
expected_classes = {
|
|
_CLASS_DOSAGE,
|
|
_CLASS_ROUTE,
|
|
_CLASS_MEDICATION,
|
|
_CLASS_FREQUENCY,
|
|
_CLASS_DURATION,
|
|
}
|
|
assert_extractions_contain(self, result, expected_classes)
|
|
assert_valid_char_intervals(self, result)
|
|
|
|
# Using regex for precise matching to avoid false positives
|
|
medication_texts = extract_by_class(result, _CLASS_MEDICATION)
|
|
self.assertTrue(
|
|
any(
|
|
re.search(r"\bIbuprofen\b", text, re.IGNORECASE)
|
|
for text in medication_texts
|
|
),
|
|
f"No Ibuprofen found in: {medication_texts}",
|
|
)
|
|
|
|
dosage_texts = extract_by_class(result, _CLASS_DOSAGE)
|
|
self.assertTrue(
|
|
any(
|
|
re.search(r"\b400\s*mg\b", text, re.IGNORECASE)
|
|
for text in dosage_texts
|
|
),
|
|
f"No 400mg dosage found in: {dosage_texts}",
|
|
)
|
|
|
|
route_texts = extract_by_class(result, _CLASS_ROUTE)
|
|
self.assertTrue(
|
|
any(
|
|
re.search(r"\b(PO|oral)\b", text, re.IGNORECASE)
|
|
for text in route_texts
|
|
),
|
|
f"No PO/oral route found in: {route_texts}",
|
|
)
|
|
|
|
@skip_if_no_gemini
|
|
@live_api
|
|
@retry_on_transient_errors(max_retries=2)
|
|
def test_multilingual_medication_extraction(self):
|
|
"""Test medication extraction with Japanese text."""
|
|
text = ( # "The patient takes 10 mg of medication daily."
|
|
"患者は毎日10mgの薬を服用します。"
|
|
)
|
|
|
|
prompt = "Extract medication information including dosage and frequency."
|
|
|
|
examples = [
|
|
lx.data.ExampleData(
|
|
text="The patient takes 20mg of aspirin twice daily.",
|
|
extractions=[
|
|
lx.data.Extraction(
|
|
extraction_class=_CLASS_MEDICATION,
|
|
extraction_text="aspirin",
|
|
attributes={
|
|
_CLASS_DOSAGE: "20mg",
|
|
_CLASS_FREQUENCY: "twice daily",
|
|
},
|
|
),
|
|
],
|
|
)
|
|
]
|
|
|
|
unicode_tokenizer = tokenizer_lib.UnicodeTokenizer()
|
|
|
|
result = lx.extract(
|
|
text_or_documents=text,
|
|
prompt_description=prompt,
|
|
examples=examples,
|
|
model_id=DEFAULT_GEMINI_MODEL,
|
|
api_key=GEMINI_API_KEY,
|
|
language_model_params=GEMINI_MODEL_PARAMS,
|
|
tokenizer=unicode_tokenizer,
|
|
)
|
|
|
|
assert result is not None
|
|
self.assertIsInstance(result, lx.data.AnnotatedDocument)
|
|
assert len(result.extractions) > 0
|
|
|
|
medication_extractions = [
|
|
e for e in result.extractions if e.extraction_class == _CLASS_MEDICATION
|
|
]
|
|
assert (
|
|
len(medication_extractions) > 0
|
|
), "No medication entities found in Japanese text"
|
|
assert_valid_char_intervals(self, result)
|
|
|
|
@skip_if_no_gemini
|
|
@live_api
|
|
@retry_on_transient_errors(max_retries=2)
|
|
def test_explicit_provider_gemini(self):
|
|
"""Test using explicit provider with Gemini."""
|
|
config = lx.factory.ModelConfig(
|
|
model_id=DEFAULT_GEMINI_MODEL,
|
|
provider="GeminiLanguageModel",
|
|
provider_kwargs={
|
|
"api_key": GEMINI_API_KEY,
|
|
"temperature": 0.0,
|
|
},
|
|
)
|
|
|
|
model = lx.factory.create_model(config)
|
|
self.assertEqual(model.__class__.__name__, "GeminiLanguageModel")
|
|
self.assertEqual(model.model_id, DEFAULT_GEMINI_MODEL)
|
|
|
|
config2 = lx.factory.ModelConfig(
|
|
model_id=DEFAULT_GEMINI_MODEL,
|
|
provider="gemini",
|
|
provider_kwargs={
|
|
"api_key": GEMINI_API_KEY,
|
|
},
|
|
)
|
|
|
|
model2 = lx.factory.create_model(config2)
|
|
self.assertEqual(model2.__class__.__name__, "GeminiLanguageModel")
|
|
|
|
@skip_if_no_gemini
|
|
@live_api
|
|
@retry_on_transient_errors(max_retries=2)
|
|
def test_medication_relationship_extraction(self):
|
|
"""Test relationship extraction for medications with Gemini."""
|
|
input_text = """
|
|
The patient was prescribed Lisinopril and Metformin last month.
|
|
He takes the Lisinopril 10mg daily for hypertension, but often misses
|
|
his Metformin 500mg dose which should be taken twice daily for diabetes.
|
|
"""
|
|
|
|
prompt = textwrap.dedent("""
|
|
Extract medications with their details, using attributes to group related information:
|
|
|
|
1. Extract entities in the order they appear in the text
|
|
2. Each entity must have a 'medication_group' attribute linking it to its medication
|
|
3. All details about a medication should share the same medication_group value
|
|
""")
|
|
|
|
examples = get_relationship_examples()
|
|
|
|
result = lx.extract(
|
|
text_or_documents=input_text,
|
|
prompt_description=prompt,
|
|
examples=examples,
|
|
model_id=DEFAULT_GEMINI_MODEL,
|
|
api_key=GEMINI_API_KEY,
|
|
language_model_params=GEMINI_MODEL_PARAMS,
|
|
)
|
|
|
|
assert result is not None
|
|
assert len(result.extractions) > 0
|
|
assert_valid_char_intervals(self, result)
|
|
|
|
medication_groups = {}
|
|
for extraction in result.extractions:
|
|
assert (
|
|
extraction.attributes is not None
|
|
), f"Missing attributes for {extraction.extraction_text}"
|
|
assert (
|
|
"medication_group" in extraction.attributes
|
|
), f"Missing medication_group for {extraction.extraction_text}"
|
|
|
|
group_name = extraction.attributes["medication_group"]
|
|
medication_groups.setdefault(group_name, []).append(extraction)
|
|
|
|
assert (
|
|
len(medication_groups) >= 2
|
|
), f"Expected at least 2 medications, found {len(medication_groups)}"
|
|
|
|
# Allow flexible matching for dosage field (could be "dosage" or "dose")
|
|
for med_name, extractions in medication_groups.items():
|
|
extraction_classes = {e.extraction_class for e in extractions}
|
|
# At minimum, each group should have the medication itself
|
|
assert (
|
|
_CLASS_MEDICATION in extraction_classes
|
|
), f"{med_name} group missing medication entity"
|
|
# Dosage is expected but might be formatted differently
|
|
assert any(
|
|
c in extraction_classes for c in [_CLASS_DOSAGE, "dose"]
|
|
), f"{med_name} group missing dosage"
|
|
|
|
@skip_if_no_vertex
|
|
@live_api
|
|
@pytest.mark.vertex_ai
|
|
@mock.patch.object(gb, "infer_batch", wraps=gb.infer_batch, autospec=True)
|
|
def test_batch_extraction_vertex_gcs(self, mock_infer_batch):
|
|
"""Test extraction using Vertex AI Batch API with GCS.
|
|
|
|
This test runs a real Vertex AI Batch job and will take time to complete.
|
|
It is skipped unless VERTEX_PROJECT is set.
|
|
|
|
We wrap `infer_batch` to verify that:
|
|
- Batch API is actually called (not falling back to real-time API)
|
|
- Schema dict is passed (non-None) to the batch function
|
|
"""
|
|
|
|
prompt = textwrap.dedent("""\
|
|
Extract medication information including medication name, dosage, route, frequency,
|
|
and duration in the order they appear in the text.""")
|
|
|
|
examples = get_basic_medication_examples()
|
|
|
|
documents = [
|
|
lx.data.Document(
|
|
document_id="vx_doc1",
|
|
text="Patient took 400 mg PO Ibuprofen q4h for two days.",
|
|
),
|
|
lx.data.Document(
|
|
document_id="vx_doc2",
|
|
text="Patient was given 250 mg IV Cefazolin TID for one week.",
|
|
),
|
|
lx.data.Document(
|
|
document_id="vx_doc3",
|
|
text="Administered 2 mg IV Morphine once for acute pain.",
|
|
),
|
|
lx.data.Document(
|
|
document_id="vx_doc4",
|
|
text="Prescribed 500 mg PO Amoxicillin BID for infection.",
|
|
),
|
|
lx.data.Document(
|
|
document_id="vx_doc5",
|
|
text="Given 10 mg IM Haloperidol PRN for agitation.",
|
|
),
|
|
]
|
|
expected_meds = [
|
|
"Ibuprofen",
|
|
"Cefazolin",
|
|
"Morphine",
|
|
"Amoxicillin",
|
|
"Haloperidol",
|
|
]
|
|
|
|
language_model_params = dict(GEMINI_MODEL_PARAMS)
|
|
language_model_params["vertexai"] = True
|
|
language_model_params["project"] = VERTEX_PROJECT
|
|
language_model_params["location"] = VERTEX_LOCATION
|
|
language_model_params["batch"] = {
|
|
"enabled": True,
|
|
"threshold": 2,
|
|
"poll_interval": 1, # Fast polling for test
|
|
"timeout": 900, # 15 minutes for actual batch job completion
|
|
}
|
|
|
|
batch_result = lx.extract(
|
|
text_or_documents=documents,
|
|
prompt_description=prompt,
|
|
examples=examples,
|
|
model_id=DEFAULT_GEMINI_MODEL,
|
|
language_model_params=language_model_params,
|
|
)
|
|
|
|
mock_infer_batch.assert_called_once()
|
|
call_args = mock_infer_batch.call_args
|
|
schema_dict_arg = call_args.kwargs.get("schema_dict")
|
|
self.assertIsNotNone(
|
|
schema_dict_arg,
|
|
"schema_dict should be passed to batch API (not None)",
|
|
)
|
|
|
|
self.assertIsInstance(batch_result, list)
|
|
self.assertEqual(
|
|
len(batch_result),
|
|
len(documents),
|
|
f"Expected {len(documents)} results from Vertex batch API",
|
|
)
|
|
|
|
for i, (res, med_name) in enumerate(zip(batch_result, expected_meds)):
|
|
self.assertIsInstance(
|
|
res,
|
|
lx.data.AnnotatedDocument,
|
|
f"Result {i} should be an AnnotatedDocument, got {type(res)}",
|
|
)
|
|
self.assertTrue(
|
|
res.extractions,
|
|
f"No extractions for document {i}",
|
|
)
|
|
for extraction in res.extractions:
|
|
self.assertIsInstance(
|
|
extraction,
|
|
lx.data.Extraction,
|
|
"Extraction item should be Extraction object, got"
|
|
f" {type(extraction)}",
|
|
)
|
|
|
|
meds = extract_by_class(res, _CLASS_MEDICATION)
|
|
self.assertTrue(
|
|
any(
|
|
re.search(rf"\b{re.escape(med_name)}\b", m, re.IGNORECASE)
|
|
for m in meds
|
|
),
|
|
f"Expected medication '{med_name}' not found in results: {meds}",
|
|
)
|
|
|
|
dosages = extract_by_class(res, _CLASS_DOSAGE)
|
|
self.assertTrue(
|
|
dosages,
|
|
f"No dosage extracted for medication '{med_name}'",
|
|
)
|
|
|
|
assert_valid_char_intervals(self, res)
|
|
|
|
@skip_if_no_vertex
|
|
@live_api
|
|
@pytest.mark.vertex_ai
|
|
def test_batch_caching_live(self):
|
|
"""Test batch caching with real Vertex AI Batch API.
|
|
|
|
Verifies that:
|
|
1. First run populates GCS cache
|
|
2. Second run uses cache (returns same results faster)
|
|
"""
|
|
prompt = "Extract the medication: Patient takes 10mg Lisinopril."
|
|
examples = get_basic_medication_examples()
|
|
|
|
# Use unique IDs to ensure cache isolation between test runs.
|
|
run_id = uuid.uuid4().hex[:8]
|
|
documents = [
|
|
lx.data.Document(
|
|
document_id=f"doc_{i}_{run_id}",
|
|
text=f"Patient takes 10mg Lisinopril {i} {run_id}.",
|
|
)
|
|
for i in range(2)
|
|
]
|
|
|
|
language_model_params = dict(GEMINI_MODEL_PARAMS)
|
|
language_model_params["vertexai"] = True
|
|
language_model_params["project"] = VERTEX_PROJECT
|
|
language_model_params["location"] = VERTEX_LOCATION
|
|
language_model_params["batch"] = {
|
|
"enabled": True,
|
|
"threshold": 2,
|
|
"poll_interval": 1,
|
|
"timeout": 900,
|
|
"enable_caching": True,
|
|
}
|
|
|
|
print("\nStarting first batch run (API)...")
|
|
start_time = time.time()
|
|
results1 = list(
|
|
lx.extract(
|
|
text_or_documents=documents,
|
|
prompt_description=prompt,
|
|
examples=examples,
|
|
model_id=DEFAULT_GEMINI_MODEL,
|
|
language_model_params=language_model_params,
|
|
)
|
|
)
|
|
duration1 = time.time() - start_time
|
|
print(f"First run took {duration1:.2f}s")
|
|
|
|
print("Starting second batch run (Cache)...")
|
|
start_time = time.time()
|
|
results2 = list(
|
|
lx.extract(
|
|
text_or_documents=documents,
|
|
prompt_description=prompt,
|
|
examples=examples,
|
|
model_id=DEFAULT_GEMINI_MODEL,
|
|
language_model_params=language_model_params,
|
|
)
|
|
)
|
|
duration2 = time.time() - start_time
|
|
print(f"Second run took {duration2:.2f}s")
|
|
|
|
self.assertEqual(len(results1), len(results2))
|
|
for r1, r2 in zip(results1, results2):
|
|
self.assertEqual(r1.text, r2.text)
|
|
self.assertEqual(len(r1.extractions), len(r2.extractions))
|
|
|
|
self.assertLess(duration2, 10.0, "Second run took too long for cache hit")
|
|
|
|
print("\nVerifying GCS cache content...")
|
|
bucket_name = gb._get_bucket_name(VERTEX_PROJECT, VERTEX_LOCATION)
|
|
print(f"Checking bucket: {bucket_name}")
|
|
self._verify_gcs_cache_content(bucket_name)
|
|
|
|
|
|
class TestCrossChunkContext(unittest.TestCase):
|
|
"""Tests for cross-chunk context feature with real API."""
|
|
|
|
@skip_if_no_gemini
|
|
@live_api
|
|
@retry_on_transient_errors(max_retries=3)
|
|
def test_context_window_extracts_from_both_chunks(self):
|
|
"""Verify context_window_chars enables extraction across chunk boundaries."""
|
|
input_text = (
|
|
"Dr. Sarah Chen is the lead researcher at the institute. "
|
|
"She published groundbreaking work on neural networks last year."
|
|
)
|
|
prompt = textwrap.dedent(
|
|
"""\
|
|
Extract all person names, roles, and achievements mentioned in the text.
|
|
Include both explicit names and information associated with pronouns."""
|
|
)
|
|
examples = [
|
|
lx.data.ExampleData(
|
|
text=(
|
|
"Professor James Miller leads the physics department. "
|
|
"He won the Nobel Prize in 2020."
|
|
),
|
|
extractions=[
|
|
lx.data.Extraction(
|
|
extraction_class="person",
|
|
extraction_text="Professor James Miller",
|
|
attributes={"role": "leads the physics department"},
|
|
),
|
|
lx.data.Extraction(
|
|
extraction_class="achievement",
|
|
extraction_text="won the Nobel Prize in 2020",
|
|
),
|
|
],
|
|
)
|
|
]
|
|
|
|
result = lx.extract(
|
|
text_or_documents=input_text,
|
|
prompt_description=prompt,
|
|
examples=examples,
|
|
model_id=DEFAULT_GEMINI_MODEL,
|
|
api_key=GEMINI_API_KEY,
|
|
language_model_params=GEMINI_MODEL_PARAMS,
|
|
max_char_buffer=60,
|
|
context_window_chars=50,
|
|
)
|
|
|
|
self.assertIsNotNone(result)
|
|
self.assertGreater(len(result.extractions), 0)
|
|
|
|
all_extraction_text = " ".join(
|
|
str(e.extraction_text) + " " + str(e.attributes)
|
|
for e in result.extractions
|
|
).lower()
|
|
|
|
has_chunk1_content = any(
|
|
term in all_extraction_text
|
|
for term in ("sarah", "chen", "researcher", "lead")
|
|
)
|
|
has_chunk2_content = any(
|
|
term in all_extraction_text
|
|
for term in ("published", "groundbreaking", "neural", "networks")
|
|
)
|
|
|
|
self.assertTrue(
|
|
has_chunk1_content,
|
|
f"Expected chunk 1 content (Sarah Chen). Got: {result.extractions}",
|
|
)
|
|
self.assertTrue(
|
|
has_chunk2_content,
|
|
f"Expected chunk 2 content (publication). Got: {result.extractions}",
|
|
)
|
|
|
|
|
|
class TestLiveAPIOpenAI(unittest.TestCase):
|
|
"""Tests using real OpenAI API."""
|
|
|
|
@skip_if_no_openai
|
|
@skip_if_openai_batch_live_disabled
|
|
@live_api
|
|
@retry_on_transient_errors(max_retries=1)
|
|
@mock.patch.object(
|
|
openai_batch, "infer_batch", wraps=openai_batch.infer_batch, autospec=True
|
|
)
|
|
def test_batch_extraction_uses_openai_batch_api(self, mock_infer_batch):
|
|
"""OpenAI batch mode runs a real Batch API extraction."""
|
|
prompt = textwrap.dedent("""\
|
|
Extract medication information including medication name, dosage, route,
|
|
frequency, and duration in the order they appear in the text.""")
|
|
examples = get_basic_medication_examples()
|
|
documents = [
|
|
lx.data.Document(
|
|
document_id="openai_batch_doc1",
|
|
text="Patient took 400 mg PO Ibuprofen q4h for two days.",
|
|
),
|
|
lx.data.Document(
|
|
document_id="openai_batch_doc2",
|
|
text="Administered 2 mg IV Morphine once for acute pain.",
|
|
),
|
|
]
|
|
expected_meds = ["Ibuprofen", "Morphine"]
|
|
language_model_params = {
|
|
**OPENAI_MODEL_PARAMS,
|
|
"max_output_tokens": 512,
|
|
"batch": {
|
|
"enabled": True,
|
|
"threshold": 1,
|
|
"poll_interval": 5,
|
|
"timeout": 900,
|
|
},
|
|
}
|
|
|
|
batch_result = lx.extract(
|
|
text_or_documents=documents,
|
|
prompt_description=prompt,
|
|
examples=examples,
|
|
model_id="gpt-4o-mini",
|
|
api_key=OPENAI_API_KEY,
|
|
use_schema_constraints=False,
|
|
language_model_params=language_model_params,
|
|
)
|
|
|
|
mock_infer_batch.assert_called()
|
|
call_args = mock_infer_batch.call_args
|
|
self.assertTrue(call_args.kwargs["cfg"].enabled)
|
|
self.assertEqual(call_args.kwargs["cfg"].threshold, 1)
|
|
|
|
self.assertIsInstance(batch_result, list)
|
|
self.assertEqual(len(batch_result), len(documents))
|
|
for result, expected_med in zip(batch_result, expected_meds):
|
|
self.assertIsInstance(result, lx.data.AnnotatedDocument)
|
|
medication_texts = extract_by_class(result, _CLASS_MEDICATION)
|
|
self.assertIn(expected_med, medication_texts)
|
|
assert_valid_char_intervals(self, result)
|
|
|
|
@skip_if_no_openai
|
|
@live_api
|
|
@retry_on_transient_errors(max_retries=2)
|
|
def test_medication_extraction(self):
|
|
"""Test medication extraction with OpenAI models."""
|
|
prompt = textwrap.dedent("""\
|
|
Extract medication information including medication name, dosage, route, frequency,
|
|
and duration in the order they appear in the text.""")
|
|
|
|
examples = get_basic_medication_examples()
|
|
input_text = "Patient took 400 mg PO Ibuprofen q4h for two days."
|
|
|
|
result = lx.extract(
|
|
text_or_documents=input_text,
|
|
prompt_description=prompt,
|
|
examples=examples,
|
|
model_id=DEFAULT_OPENAI_MODEL,
|
|
api_key=OPENAI_API_KEY,
|
|
use_schema_constraints=False,
|
|
language_model_params=OPENAI_MODEL_PARAMS,
|
|
)
|
|
|
|
assert result is not None
|
|
self.assertIsInstance(result, lx.data.AnnotatedDocument)
|
|
assert len(result.extractions) > 0
|
|
|
|
expected_classes = {
|
|
_CLASS_DOSAGE,
|
|
_CLASS_ROUTE,
|
|
_CLASS_MEDICATION,
|
|
_CLASS_FREQUENCY,
|
|
_CLASS_DURATION,
|
|
}
|
|
assert_extractions_contain(self, result, expected_classes)
|
|
assert_valid_char_intervals(self, result)
|
|
|
|
# Using regex for precise matching to avoid false positives
|
|
medication_texts = extract_by_class(result, _CLASS_MEDICATION)
|
|
self.assertTrue(
|
|
any(
|
|
re.search(r"\bIbuprofen\b", text, re.IGNORECASE)
|
|
for text in medication_texts
|
|
),
|
|
f"No Ibuprofen found in: {medication_texts}",
|
|
)
|
|
|
|
dosage_texts = extract_by_class(result, _CLASS_DOSAGE)
|
|
self.assertTrue(
|
|
any(
|
|
re.search(r"\b400\s*mg\b", text, re.IGNORECASE)
|
|
for text in dosage_texts
|
|
),
|
|
f"No 400mg dosage found in: {dosage_texts}",
|
|
)
|
|
|
|
route_texts = extract_by_class(result, _CLASS_ROUTE)
|
|
self.assertTrue(
|
|
any(
|
|
re.search(r"\b(PO|oral)\b", text, re.IGNORECASE)
|
|
for text in route_texts
|
|
),
|
|
f"No PO/oral route found in: {route_texts}",
|
|
)
|
|
|
|
@skip_if_no_openai
|
|
@live_api
|
|
@retry_on_transient_errors(max_retries=2)
|
|
def test_medication_extraction_with_schema_constraints(self):
|
|
"""Strict OpenAI outputs enforce the example-derived extraction shape."""
|
|
prompt = textwrap.dedent("""\
|
|
Extract conditions and medications in the order they appear in the text.
|
|
Use exact text for extractions. For condition attributes, include status
|
|
and symptoms as a list when symptoms are available.""")
|
|
examples = [
|
|
lx.data.ExampleData(
|
|
text="Patient has diabetes with fatigue and takes Metformin.",
|
|
extractions=[
|
|
lx.data.Extraction(
|
|
extraction_class=_CLASS_CONDITION,
|
|
extraction_text="diabetes",
|
|
attributes={
|
|
"status": "present",
|
|
"symptoms": ["fatigue"],
|
|
},
|
|
),
|
|
lx.data.Extraction(
|
|
extraction_class=_CLASS_MEDICATION,
|
|
extraction_text="Metformin",
|
|
attributes={"status": "current"},
|
|
),
|
|
],
|
|
)
|
|
]
|
|
input_text = (
|
|
"Patient has headache with fatigue and chills and took 400 mg PO "
|
|
"Ibuprofen."
|
|
)
|
|
|
|
result = lx.extract(
|
|
text_or_documents=input_text,
|
|
prompt_description=prompt,
|
|
examples=examples,
|
|
model_id="gpt-4o-mini",
|
|
api_key=OPENAI_API_KEY,
|
|
use_schema_constraints=True,
|
|
fence_output=False,
|
|
language_model_params={
|
|
**OPENAI_MODEL_PARAMS,
|
|
"max_output_tokens": 512,
|
|
},
|
|
)
|
|
|
|
self.assertIsInstance(result, lx.data.AnnotatedDocument)
|
|
self.assertGreater(len(result.extractions), 0)
|
|
allowed_classes = {_CLASS_CONDITION, _CLASS_MEDICATION}
|
|
extraction_classes = {
|
|
extraction.extraction_class for extraction in result.extractions
|
|
}
|
|
self.assertSetEqual(extraction_classes, allowed_classes)
|
|
allowed_attribute_keys = {
|
|
_CLASS_CONDITION: {"status", "symptoms"},
|
|
_CLASS_MEDICATION: {"status"},
|
|
}
|
|
for extraction in result.extractions:
|
|
if isinstance(extraction.attributes, dict):
|
|
self.assertLessEqual(
|
|
set(extraction.attributes),
|
|
allowed_attribute_keys[extraction.extraction_class],
|
|
)
|
|
condition_extractions = [
|
|
extraction
|
|
for extraction in result.extractions
|
|
if extraction.extraction_class == _CLASS_CONDITION
|
|
]
|
|
self.assertTrue(
|
|
any(
|
|
isinstance(extraction.attributes, dict)
|
|
and isinstance(extraction.attributes.get("symptoms"), list)
|
|
for extraction in condition_extractions
|
|
),
|
|
f"Expected list-valued symptoms attribute. Got: {result.extractions}",
|
|
)
|
|
assert_valid_char_intervals(self, result)
|
|
|
|
@skip_if_no_openai
|
|
@live_api
|
|
@retry_on_transient_errors(max_retries=2)
|
|
def test_explicit_provider_selection(self):
|
|
"""Test using explicit provider parameter for disambiguation."""
|
|
# Test with explicit model_id and provider
|
|
config = lx.factory.ModelConfig(
|
|
model_id=DEFAULT_OPENAI_MODEL,
|
|
provider="OpenAILanguageModel", # Explicit provider selection
|
|
provider_kwargs={
|
|
"api_key": OPENAI_API_KEY,
|
|
"temperature": 0.0,
|
|
},
|
|
)
|
|
|
|
model = lx.factory.create_model(config)
|
|
|
|
self.assertIsInstance(model, lx.providers.openai.OpenAILanguageModel)
|
|
self.assertEqual(model.model_id, DEFAULT_OPENAI_MODEL)
|
|
|
|
# Also test using provider without model_id (uses default)
|
|
config_default = lx.factory.ModelConfig(
|
|
provider="OpenAILanguageModel",
|
|
provider_kwargs={
|
|
"api_key": OPENAI_API_KEY,
|
|
},
|
|
)
|
|
|
|
model_default = lx.factory.create_model(config_default)
|
|
self.assertEqual(model_default.__class__.__name__, "OpenAILanguageModel")
|
|
# Should use the default model_id from the provider
|
|
self.assertEqual(model_default.model_id, "gpt-4o-mini")
|
|
|
|
@skip_if_no_openai
|
|
@live_api
|
|
@retry_on_transient_errors(max_retries=2)
|
|
def test_medication_relationship_extraction(self):
|
|
"""Test relationship extraction for medications with OpenAI."""
|
|
input_text = """
|
|
The patient was prescribed Lisinopril and Metformin last month.
|
|
He takes the Lisinopril 10mg daily for hypertension, but often misses
|
|
his Metformin 500mg dose which should be taken twice daily for diabetes.
|
|
"""
|
|
|
|
prompt = textwrap.dedent("""
|
|
Extract medications with their details, using attributes to group related information:
|
|
|
|
1. Extract entities in the order they appear in the text
|
|
2. Each entity must have a 'medication_group' attribute linking it to its medication
|
|
3. All details about a medication should share the same medication_group value
|
|
""")
|
|
|
|
examples = get_relationship_examples()
|
|
|
|
result = lx.extract(
|
|
text_or_documents=input_text,
|
|
prompt_description=prompt,
|
|
examples=examples,
|
|
model_id=DEFAULT_OPENAI_MODEL,
|
|
api_key=OPENAI_API_KEY,
|
|
use_schema_constraints=False,
|
|
language_model_params=OPENAI_MODEL_PARAMS,
|
|
)
|
|
|
|
assert result is not None
|
|
assert len(result.extractions) > 0
|
|
assert_valid_char_intervals(self, result)
|
|
|
|
medication_groups = {}
|
|
for extraction in result.extractions:
|
|
assert (
|
|
extraction.attributes is not None
|
|
), f"Missing attributes for {extraction.extraction_text}"
|
|
assert (
|
|
"medication_group" in extraction.attributes
|
|
), f"Missing medication_group for {extraction.extraction_text}"
|
|
|
|
group_name = extraction.attributes["medication_group"]
|
|
medication_groups.setdefault(group_name, []).append(extraction)
|
|
|
|
assert (
|
|
len(medication_groups) >= 2
|
|
), f"Expected at least 2 medications, found {len(medication_groups)}"
|
|
|
|
# Allow flexible matching for dosage field (could be "dosage" or "dose")
|
|
for med_name, extractions in medication_groups.items():
|
|
extraction_classes = {e.extraction_class for e in extractions}
|
|
# At minimum, each group should have the medication itself
|
|
assert (
|
|
_CLASS_MEDICATION in extraction_classes
|
|
), f"{med_name} group missing medication entity"
|
|
# Dosage is expected but might be formatted differently
|
|
assert any(
|
|
c in extraction_classes for c in [_CLASS_DOSAGE, "dose"]
|
|
), f"{med_name} group missing dosage"
|
|
|
|
@skip_if_no_openai
|
|
@live_api
|
|
@retry_on_transient_errors(max_retries=2)
|
|
def test_reasoning_effort_passthrough(self):
|
|
"""reasoning_effort is accepted by reasoning models."""
|
|
examples = get_basic_medication_examples()
|
|
input_text = "Patient took 400 mg PO Ibuprofen q4h for two days."
|
|
|
|
config = lx.factory.ModelConfig(
|
|
model_id="o4-mini",
|
|
provider="OpenAILanguageModel",
|
|
provider_kwargs={
|
|
"api_key": OPENAI_API_KEY,
|
|
"reasoning_effort": "low",
|
|
},
|
|
)
|
|
|
|
result = lx.extract(
|
|
text_or_documents=input_text,
|
|
prompt_description="Extract medications.",
|
|
examples=examples,
|
|
config=config,
|
|
use_schema_constraints=False,
|
|
)
|
|
|
|
assert result is not None
|
|
self.assertIsInstance(result, lx.data.AnnotatedDocument)
|
|
|
|
|
|
class TestLiveAPIOutputSchema(unittest.TestCase):
|
|
"""Live tests for user-provided output_schema support."""
|
|
|
|
_OUTPUT_SCHEMA = {
|
|
"type": "object",
|
|
"properties": {
|
|
"extractions": {
|
|
"type": "array",
|
|
"items": {
|
|
"type": "object",
|
|
"properties": {
|
|
"condition": {"type": "string"},
|
|
"condition_attributes": {
|
|
"type": "object",
|
|
"properties": {
|
|
"status": {
|
|
"type": "string",
|
|
"enum": ["active", "resolved"],
|
|
}
|
|
},
|
|
"required": ["status"],
|
|
"additionalProperties": False,
|
|
},
|
|
},
|
|
"required": ["condition", "condition_attributes"],
|
|
"additionalProperties": False,
|
|
},
|
|
}
|
|
},
|
|
"required": ["extractions"],
|
|
"additionalProperties": False,
|
|
}
|
|
_INPUT_TEXT = "Patient has active hypertension and a resolved infection."
|
|
|
|
def _assert_schema_constrained_extractions(self, result):
|
|
self.assertIsInstance(result, lx.data.AnnotatedDocument)
|
|
self.assertTrue(result.extractions)
|
|
for extraction in result.extractions:
|
|
self.assertEqual(extraction.extraction_class, "condition")
|
|
if extraction.attributes:
|
|
self.assertIn(
|
|
extraction.attributes.get("status"), ("active", "resolved")
|
|
)
|
|
|
|
@skip_if_no_gemini
|
|
@live_api
|
|
def test_gemini_extract_with_output_schema(self):
|
|
result = lx.extract(
|
|
text_or_documents=self._INPUT_TEXT,
|
|
prompt_description=(
|
|
"Extract medical conditions with their status attribute."
|
|
),
|
|
model_id=DEFAULT_GEMINI_MODEL,
|
|
api_key=GEMINI_API_KEY,
|
|
output_schema=self._OUTPUT_SCHEMA,
|
|
)
|
|
|
|
self._assert_schema_constrained_extractions(result)
|
|
|
|
@skip_if_no_openai
|
|
@live_api
|
|
def test_openai_extract_with_output_schema(self):
|
|
result = lx.extract(
|
|
text_or_documents=self._INPUT_TEXT,
|
|
prompt_description=(
|
|
"Extract medical conditions with their status attribute."
|
|
),
|
|
model_id=DEFAULT_OPENAI_MODEL,
|
|
api_key=OPENAI_API_KEY,
|
|
output_schema=self._OUTPUT_SCHEMA,
|
|
)
|
|
|
|
self._assert_schema_constrained_extractions(result)
|