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.core.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.core.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 effective_entity_mode(mode: WikiEntityMode) -> WikiEntityMode: 

39 """The mode that will actually run for *mode*, applying the fallback. 

40 

41 Unimplemented strategies resolve to ``NER_ENTITIES``; provenance records this 

42 so the audit reflects the extractor that ran, not the configured request. 

43 """ 

44 return mode if mode in _IMPLEMENTED_MODES else WikiEntityMode.NER_ENTITIES 

45 

46 

47def get_entity_extractor( 

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

49) -> EntityExtractor: 

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

51 

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

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

54 build or sync mid-flight. 

55 """ 

56 if mode not in _IMPLEMENTED_MODES: 

57 log.warning( 

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

59 mode.value, 

60 WikiEntityMode.NER_ENTITIES.value, 

61 ) 

62 effective = effective_entity_mode(mode) 

63 factory = _EXTRACTOR_BY_MODE[effective] 

64 return factory(provider, config)