semantica/tests/test_ner_configurations.py

import unittest
import sys
import os
from unittest.mock import MagicMock, patch
from dataclasses import asdict

# Add project root to path
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from semantica.semantic_extract.ner_extractor import NERExtractor, Entity
from semantica.semantic_extract.named_entity_recognizer import NamedEntityRecognizer
from semantica.semantic_extract.methods import get_entity_method

class TestNERConfigurations(unittest.TestCase):
    """
    Test suite to verify NER with different configurations:
    - LLM
    - ML (spaCy)
    - Regex
    - Pattern
    - Fallbacks and Ensemble
    """

    def setUp(self):
        self.text = "Apple Inc. was founded by Steve Jobs."

    @patch('semantica.semantic_extract.methods.create_provider')
    def test_ner_llm_config(self, mock_create_provider):
        """Test NER with LLM configuration"""
        print("\nTesting NER with LLM configuration...")
        
        # Mock LLM provider
        mock_provider = MagicMock()
        mock_provider.is_available.return_value = True
        mock_provider.generate_structured.return_value = [
            {"text": "Apple Inc.", "label": "ORG", "start": 0, "end": 10, "confidence": 0.95},
            {"text": "Steve Jobs", "label": "PERSON", "start": 26, "end": 36, "confidence": 0.98}
        ]
        mock_create_provider.return_value = mock_provider
        
        # Initialize extractor with LLM method
        extractor = NERExtractor(
            method="llm", 
            provider="openai", 
            model="gpt-4",
            temperature=0.1
        )
        
        entities = extractor.extract_entities(self.text)
        
        # Verify provider creation args
        mock_create_provider.assert_called_with("openai", model="gpt-4", temperature=0.1)
        
        # Verify extraction
        self.assertEqual(len(entities), 2)
        self.assertEqual(entities[0].text, "Apple Inc.")
        self.assertEqual(entities[0].label, "ORG")
        self.assertEqual(entities[0].metadata["extraction_method"], "llm")
        self.assertEqual(entities[0].metadata["model"], "gpt-4")

    @patch('semantica.semantic_extract.methods.spacy')
    def test_ner_ml_config_spacy_available(self, mock_spacy):
        """Test NER with ML (spaCy) configuration when spaCy is available"""
        print("\nTesting NER with ML (spaCy) configuration...")
        
        # Mock spaCy nlp model
        mock_nlp = MagicMock()
        mock_doc = MagicMock()
        
        # Mock entities
        ent1 = MagicMock()
        ent1.text = "Apple Inc."
        ent1.label_ = "ORG"
        ent1.start_char = 0
        ent1.end_char = 10
        ent1.confidence = 1.0 # Optional attribute
        
        ent2 = MagicMock()
        ent2.text = "Steve Jobs"
        ent2.label_ = "PERSON"
        ent2.start_char = 26
        ent2.end_char = 36
        ent2.confidence = 0.99
        
        mock_doc.ents = [ent1, ent2]
        mock_nlp.return_value = mock_doc
        mock_spacy.load.return_value = mock_nlp
        
        # Patch SPACY_AVAILABLE in methods module
        with patch('semantica.semantic_extract.methods.SPACY_AVAILABLE', True):
            extractor = NERExtractor(method="ml", model="en_core_web_trf")
            entities = extractor.extract_entities(self.text)
            
            # Verify spacy load called with correct model
            mock_spacy.load.assert_called_with("en_core_web_trf")
            
            self.assertEqual(len(entities), 2)
            self.assertEqual(entities[0].text, "Apple Inc.")
            self.assertEqual(entities[0].label, "ORG")
            self.assertEqual(entities[0].metadata["extraction_method"], "ml")
            self.assertEqual(entities[0].metadata["model"], "en_core_web_trf")

    @patch('semantica.semantic_extract.ner_extractor.spacy')
    def test_ner_ml_init_falls_back_when_spacy_runtime_is_broken(self, mock_spacy):
        """Test NER init does not crash when spaCy is installed but unusable at runtime."""
        mock_spacy.load.side_effect = RuntimeError("ConfigSchemaNlp is not fully defined")

        with patch('semantica.semantic_extract.ner_extractor.SPACY_AVAILABLE', True):
            extractor = NERExtractor(method="ml", model="en_core_web_sm")

        self.assertIsNone(extractor.nlp)
        self.assertFalse(extractor._ml_runtime_usable)

    @patch('semantica.semantic_extract.ner_extractor.spacy')
    @patch('semantica.semantic_extract.methods.get_entity_method')
    @patch('semantica.semantic_extract.methods.spacy')
    def test_ner_ml_runtime_failure_disables_repeated_ml_load_attempts(
        self,
        mock_methods_spacy,
        mock_get_method,
        mock_init_spacy,
    ):
        """Test degraded ML mode skips repeated spaCy load attempts after init failure."""
        mock_init_spacy.load.side_effect = RuntimeError("ConfigSchemaNlp is not fully defined")
        mock_ml_method = MagicMock(return_value=[])
        mock_get_method.side_effect = lambda name: mock_ml_method if name == "ml" else (lambda *_args, **_kwargs: [])

        with patch('semantica.semantic_extract.ner_extractor.SPACY_AVAILABLE', True):
            extractor = NERExtractor(method=["ml", "pattern"], model="en_core_web_sm")

        entities = extractor.extract_entities(self.text)

        self.assertFalse(extractor._ml_runtime_usable)
        self.assertEqual(mock_init_spacy.load.call_count, 1)
        self.assertEqual(mock_methods_spacy.load.call_count, 0)
        self.assertEqual(mock_ml_method.call_count, 0)
        self.assertIsInstance(entities, list)

    def test_ner_regex_config(self):
        """Test NER with Regex configuration"""
        print("\nTesting NER with Regex configuration...")
        
        custom_patterns = {
            "COMPANY": r"Apple Inc\.",
            "FOUNDER": r"Steve Jobs"
        }
        
        extractor = NERExtractor(method="regex", patterns=custom_patterns)
        entities = extractor.extract_entities(self.text)
        
        self.assertEqual(len(entities), 2)
        
        # Check if labels match custom keys
        labels = sorted([e.label for e in entities])
        self.assertEqual(labels, ["COMPANY", "FOUNDER"])
        
        # Check metadata
        self.assertEqual(entities[0].metadata["extraction_method"], "regex")

    def test_ner_pattern_config(self):
        """Test NER with default Pattern configuration"""
        print("\nTesting NER with Pattern configuration...")
        
        # Default patterns in methods.py match "Apple Inc" (ORG) and "Steve Jobs" (PERSON)
        # Note: The pattern for ORG in methods.py expects "Inc|Corp..."
        
        extractor = NERExtractor(method="pattern")
        entities = extractor.extract_entities(self.text)
        
        self.assertTrue(len(entities) >= 2)
        texts = [e.text for e in entities]
        self.assertIn("Apple Inc", texts) # Regex pattern does not capture the trailing dot
        # Actually methods.py regex: r"\b([A-Z][a-zA-Z]+(?:\s+[A-Z][a-zA-Z]+)*\s+(?:Inc|Corp|LLC|Ltd|Company))\b"
        # "Apple Inc." -> "Apple Inc" (dot is outside \b if not matched?)
        # Let's check the result strictly
        
    @patch('semantica.semantic_extract.methods.create_provider')
    @patch('semantica.semantic_extract.methods.spacy')
    def test_ner_ensemble_config(self, mock_spacy, mock_create_provider):
        """Test NER with Ensemble (Multiple Methods)"""
        print("\nTesting NER with Ensemble configuration...")
        
        # Setup mocks
        # LLM returns 1 entity
        mock_provider = MagicMock()
        mock_provider.is_available.return_value = True
        mock_provider.generate_structured.return_value = [
             {"text": "Apple Inc.", "label": "ORG", "start": 0, "end": 10, "confidence": 0.95}
        ]
        mock_create_provider.return_value = mock_provider
        
        # ML returns 2 entities
        mock_nlp = MagicMock()
        mock_doc = MagicMock()
        ent1 = MagicMock()
        ent1.text = "Apple Inc."
        ent1.label_ = "ORG"
        ent1.start_char = 0
        ent1.end_char = 10
        ent1.confidence = 0.95
        ent2 = MagicMock()
        ent2.text = "Steve Jobs"
        ent2.label_ = "PERSON"
        ent2.start_char = 26
        ent2.end_char = 36
        ent2.confidence = 0.99
        mock_doc.ents = [ent1, ent2]
        mock_nlp.return_value = mock_doc
        mock_spacy.load.return_value = mock_nlp
        
        with patch('semantica.semantic_extract.methods.SPACY_AVAILABLE', True):
            # Init extractor with list of methods
            extractor = NERExtractor(method=["llm", "ml"], ensemble_voting=True)
            entities = extractor.extract_entities(self.text)
            
            # Since ensemble_voting=True (implied merge), we expect unique entities
            # Apple Inc (from both) + Steve Jobs (from ML)
            
            texts = [e.text for e in entities]
            self.assertIn("Apple Inc.", texts)
            self.assertIn("Steve Jobs", texts)

    @patch('semantica.semantic_extract.methods.HuggingFaceModelLoader')
    def test_ner_huggingface_config(self, mock_loader_cls):
        """Test NER with HuggingFace configuration"""
        print("\nTesting NER with HuggingFace configuration...")
        
        mock_loader = MagicMock()
        mock_loader_cls.return_value = mock_loader
        
        # Mock extract_entities return
        # HuggingFace loader typically returns list of dicts or objects
        mock_loader.extract_entities.return_value = [
            {"word": "Apple Inc.", "entity_group": "ORG", "score": 0.99, "start": 0, "end": 10}
        ]
        
        extractor = NERExtractor(
            method="huggingface", 
            huggingface_model="dslim/bert-base-NER",
            device="cpu"
        )
        entities = extractor.extract_entities(self.text)
        
        mock_loader.load_ner_model.assert_called_with("dslim/bert-base-NER")
        self.assertEqual(len(entities), 1)
        self.assertEqual(entities[0].text, "Apple Inc.")
        self.assertEqual(entities[0].label, "ORG")

if __name__ == '__main__':
    unittest.main()