Coverage for core / src / sensorkit / config / core.py: 84%

1# SPDX-License-Identifier: Apache-2.0 

2from pydantic import BaseModel, model_validator 

3 

4from sensorkit.auto.agent import AgentConfig 

5from sensorkit.auto.agent import __name__ as agent_module_name 

6from sensorkit.common.keyword import KeywordDict 

7from sensorkit.config.section import declare_config_section 

8from sensorkit.core.entity import EntityRef 

9from sensorkit.data.graph import DataGraph 

10 

11declare_config_section( 

12 "automation", 

13 AgentConfig, 

14 id_source="by_key", 

15 id_default="agent", 

16 service_path=agent_module_name, 

17) 

18 

19 

20class DataFlowConfig(BaseModel): 

21 entity: EntityRef 

22 producer: DataGraph | None = None 

23 consumer: DataGraph | None = None 

24 

25 @model_validator(mode="after") 

26 def _validator(self): 

27 if (self.producer is None) == (self.consumer is None): 

28 raise ValueError("Exactly one of 'producer' or 'consumer' must be set") 

29 

30 return self 

31 

32 

33declare_config_section( 

34 "data_flow", 

35 list[DataFlowConfig], 

36 id_source="by_subkey", 

37 id_key="entity", 

38 model_mapper=lambda obj: (config.producer or config.consumer for config in obj), 

39) 

40 

41 

42# Each entity carries a whole mapping of keywords, so this section produces a model per 

43# keyword rather than one per entity, repeating the entity ID across them. 

44declare_config_section( 

45 "config", 

46 dict[str, KeywordDict], 

47 id_source="mapping_key", 

48 id_mapper=lambda raw: (entity for entity, kwdict in raw.items() for _ in kwdict), 

49 model_mapper=lambda obj: (model for kwdict in obj.values() for model in kwdict.values()), 

50)