import unittest from unittest.mock import MagicMock, patch, ANY import sys import os from typing import List, Optional from pydantic import BaseModel import importlib.util # Add project root to path sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))) # Save originals before injecting mocks so they can be restored after import. _MOCKED_MODULES = ["spacy", "instructor", "groq", "openai", "sentence_transformers", "transformers"] _original_modules = {k: sys.modules.get(k) for k in _MOCKED_MODULES} # Mock external dependencies with __spec__ for importlib checks mock_spacy = MagicMock() mock_spacy.__spec__ = MagicMock() sys.modules["spacy"] = mock_spacy sys.modules["instructor"] = MagicMock() sys.modules["groq"] = MagicMock() # Better mock for openai mock_openai = MagicMock() mock_openai.__spec__ = MagicMock() sys.modules["openai"] = mock_openai # Mock sentence_transformers and transformers to avoid heavy imports and dependency checks sys.modules["sentence_transformers"] = MagicMock() mock_transformers = MagicMock() mock_transformers.__spec__ = MagicMock() sys.modules["transformers"] = mock_transformers from semantica.semantic_extract import NERExtractor from semantica.semantic_extract.methods import extract_entities_llm, _extract_entities_chunked, extract_relations_llm, extract_triplets_llm from semantica.semantic_extract.providers import BaseProvider # Restore real sys.modules entries now that the mocks have served their purpose # for the imports above. Leaving them in place would poison other test modules # (e.g. test_pr482_deepseek_openai) that need the real packages at test run time. for _key, _original in _original_modules.items(): if _original is None: sys.modules.pop(_key, None) else: sys.modules[_key] = _original class EntitiesResponse(BaseModel): entities: List[dict] class TestRetryLogic(unittest.TestCase): def setUp(self): self.mock_provider = MagicMock() self.mock_provider.is_available.return_value = True self.mock_provider.generate_typed.return_value = MagicMock(entities=[]) # Clear the module-level extraction cache to avoid cross-test interference from semantica.semantic_extract.methods import _result_cache _result_cache.clear() def test_ner_extractor_init_default(self): """Test default max_retries in NERExtractor""" ner = NERExtractor(method="llm", provider="test") # Check internal config, max_retries not in config means default behavior downstream self.assertIsNone(ner.config.get("max_retries")) def test_ner_extractor_init_custom(self): """Test custom max_retries in NERExtractor init""" ner = NERExtractor(method="llm", provider="test", max_retries=5) self.assertEqual(ner.config.get("max_retries"), 5) @patch('semantica.semantic_extract.methods.create_provider') def test_extract_entities_uses_init_value(self, mock_create_provider): """Test extract_entities uses initialized max_retries""" mock_create_provider.return_value = self.mock_provider ner = NERExtractor(method="llm", provider="test", max_retries=5) ner.extract_entities("test text") # Verify generate_typed called with max_retries=5 args, kwargs = self.mock_provider.generate_typed.call_args self.assertEqual(kwargs.get("max_retries"), 5) @patch('semantica.semantic_extract.methods.create_provider') def test_extract_entities_override(self, mock_create_provider): """Test extract_entities override max_retries""" mock_create_provider.return_value = self.mock_provider ner = NERExtractor(method="llm", provider="test", max_retries=5) # Override with 1 ner.extract_entities("test text", max_retries=1) args, kwargs = self.mock_provider.generate_typed.call_args self.assertEqual(kwargs.get("max_retries"), 1) @patch('semantica.semantic_extract.methods.create_provider') def test_chunked_extraction_propagation(self, mock_create_provider): """Test max_retries propagation in chunked extraction""" mock_create_provider.return_value = self.mock_provider # Patch TextSplitter where it lives with patch('semantica.split.TextSplitter') as MockSplitter: mock_splitter_instance = MockSplitter.return_value # Mock split to return 2 chunks mock_chunk1 = MagicMock() mock_chunk1.text = "chunk1" mock_chunk2 = MagicMock() mock_chunk2.text = "chunk2" mock_splitter_instance.split.return_value = [mock_chunk1, mock_chunk2] # Force chunking by setting max_text_length small extract_entities_llm( "very long text", provider="test", model="test-model", max_text_length=10, # Force chunking max_retries=7, structured_output_mode="typed" ) # Check if generate_typed was called with max_retries=7 for chunks # It should be called twice (once for each chunk) self.assertEqual(self.mock_provider.generate_typed.call_count, 2) # Check arguments of the calls call_args_list = self.mock_provider.generate_typed.call_args_list for args, kwargs in call_args_list: self.assertEqual(kwargs.get("max_retries"), 7) def test_provider_base_logic(self): """Test BaseProvider logic for max_retries with manual loop""" provider = BaseProvider() provider.client = MagicMock() provider.logger = MagicMock() provider.generate_structured = MagicMock(side_effect=Exception("Fail")) # Mock instructor failing with patch('semantica.semantic_extract.providers.instructor') as mock_instructor: # Make instructor client fail mock_client = MagicMock() mock_client.chat.completions.create.side_effect = Exception("Instructor Fail") mock_instructor.from_provider.return_value = mock_client mock_instructor.from_openai.return_value = mock_client # Run with max_retries=2 try: provider.generate_typed("prompt", EntitiesResponse, max_retries=2) except Exception: pass # Should try manual generation exactly 2 times self.assertEqual(provider.generate_structured.call_count, 2) def test_provider_zero_retries(self): """Test BaseProvider with max_retries=0""" provider = BaseProvider() provider.client = MagicMock() provider.logger = MagicMock() provider.generate_structured = MagicMock(side_effect=Exception("Fail")) # Mock instructor failing with patch('semantica.semantic_extract.providers.instructor') as mock_instructor: mock_client = MagicMock() mock_client.chat.completions.create.side_effect = Exception("Instructor Fail") mock_instructor.from_provider.return_value = mock_client mock_instructor.from_openai.return_value = mock_client try: provider.generate_typed("prompt", EntitiesResponse, max_retries=0) except Exception: pass # Should NOT try manual generation loop (range(0) is empty) self.assertEqual(provider.generate_structured.call_count, 0) @patch('semantica.semantic_extract.methods.create_provider') def test_relations_retry_propagation(self, mock_create_provider): """Test max_retries propagation in relation extraction""" mock_create_provider.return_value = self.mock_provider # Create a mock entity mock_entity = MagicMock() mock_entity.text = "entity" mock_entity.start_char = 0 mock_entity.end_char = 5 extract_relations_llm( "test text", entities=[mock_entity], provider="test", max_retries=4 ) args, kwargs = self.mock_provider.generate_typed.call_args self.assertEqual(kwargs.get("max_retries"), 4) @patch('semantica.semantic_extract.methods.create_provider') def test_relations_chunked_propagation(self, mock_create_provider): """Test max_retries propagation in chunked relation extraction""" mock_create_provider.return_value = self.mock_provider with patch('semantica.split.TextSplitter') as MockSplitter: mock_splitter_instance = MockSplitter.return_value mock_chunk1 = MagicMock() mock_chunk1.text = "chunk1" mock_chunk1.start_index = 0 mock_chunk1.end_index = 6 mock_splitter_instance.split.return_value = [mock_chunk1] # Create a mock entity mock_entity = MagicMock() mock_entity.text = "entity" mock_entity.start_char = 0 mock_entity.end_char = 5 extract_relations_llm( "very long text", entities=[mock_entity], provider="test", max_text_length=10, max_retries=6 ) # Check call count - should be called for the chunk # Note: _extract_relations_chunked creates a new future for each chunk # which calls extract_relations_llm, which calls generate_typed self.assertEqual(self.mock_provider.generate_typed.call_count, 1) args, kwargs = self.mock_provider.generate_typed.call_args self.assertEqual(kwargs.get("max_retries"), 6) @patch('semantica.semantic_extract.methods.create_provider') def test_triplets_retry_propagation(self, mock_create_provider): """Test max_retries propagation in triplet extraction""" mock_create_provider.return_value = self.mock_provider extract_triplets_llm( "test text", entities=[], relations=[], provider="test", max_retries=7 ) args, kwargs = self.mock_provider.generate_typed.call_args self.assertEqual(kwargs.get("max_retries"), 7) @patch('semantica.semantic_extract.methods.create_provider') def test_triplets_chunked_propagation(self, mock_create_provider): """Test max_retries propagation in chunked triplet extraction""" mock_create_provider.return_value = self.mock_provider with patch('semantica.split.TextSplitter') as MockSplitter: mock_splitter_instance = MockSplitter.return_value mock_chunk1 = MagicMock() mock_chunk1.text = "chunk1" mock_chunk1.start_index = 0 mock_chunk1.end_index = 6 mock_splitter_instance.split.return_value = [mock_chunk1] # Use max_text_length > 100 to pass the minimum viable chunk size check # and make text longer than that extract_triplets_llm( "very long text " * 20, # length > 101 entities=[], relations=[], provider="test", max_text_length=101, max_retries=8 ) # Check call count self.assertEqual(self.mock_provider.generate_typed.call_count, 1) args, kwargs = self.mock_provider.generate_typed.call_args self.assertEqual(kwargs.get("max_retries"), 8) if __name__ == "__main__": unittest.main()