Coverage for src / lilbee / wiki / entity_extractor / factory.py: 100%

1"""Runtime selector for the entity-extraction strategy.""" 

2 

3from __future__ import annotations 

4 

5import logging 

6from collections.abc import Callable 

7from typing import TYPE_CHECKING 

8 

9from lilbee.config import WikiEntityMode 

10from lilbee.wiki.entity_extractor.base import EntityExtractor 

11from lilbee.wiki.entity_extractor.llm_tagged import LlmTaggedExtractor 

12from lilbee.wiki.entity_extractor.ner_concepts import NerConceptsExtractor 

13from lilbee.wiki.entity_extractor.ner_concepts_plus_llm_types import ( 

14 NerConceptsPlusLlmTypesExtractor, 

15) 

16 

17if TYPE_CHECKING: 

18 from lilbee.config import Config 

19 from lilbee.providers.base import LLMProvider 

20 

21log = logging.getLogger(__name__) 

22 

23_EXTRACTOR_BY_MODE: dict[ 

24 WikiEntityMode, 

25 Callable[[LLMProvider, Config], EntityExtractor], 

26] = { 

27 WikiEntityMode.NER_ENTITIES: NerConceptsExtractor, 

28 WikiEntityMode.NER_CONCEPTS_PLUS_LLM_TYPES: NerConceptsPlusLlmTypesExtractor, 

29 WikiEntityMode.LLM_TAGGED: LlmTaggedExtractor, 

30} 

31 

32# Implementations whose ``extract`` actually runs. Modes outside this set 

33# are accepted for forward compatibility (so config files and env vars 

34# keep parsing) but fall back to ``NER_ENTITIES`` with a warning. 

35_IMPLEMENTED_MODES: frozenset[WikiEntityMode] = frozenset({WikiEntityMode.NER_ENTITIES}) 

36 

37 

38def get_entity_extractor( 

39 mode: WikiEntityMode, provider: LLMProvider, config: Config 

40) -> EntityExtractor: 

41 """Return an ``EntityExtractor`` implementation for *mode*. 

42 

43 Unimplemented strategies fall back to ``NER_ENTITIES`` with a 

44 warning so a user who flips the config to a stub never crashes a 

45 build or sync mid-flight. 

46 """ 

47 if mode not in _IMPLEMENTED_MODES: 

48 log.warning( 

49 "Entity-extraction mode %r is not yet implemented; falling back to %r", 

50 mode.value, 

51 WikiEntityMode.NER_ENTITIES.value, 

52 ) 

53 mode = WikiEntityMode.NER_ENTITIES 

54 factory = _EXTRACTOR_BY_MODE[mode] 

55 return factory(provider, config)