Coverage for core / src / sensorkit / std / safety.py: 93%

1# SPDX-License-Identifier: Apache-2.0 

2from __future__ import annotations 

3 

4from typing import Literal, override 

5 

6from pydantic import BaseModel 

7 

8import sensorkit.api as sk 

9from sensorkit.auto.constraint import Constraint, ConstraintEvaluator 

10from sensorkit.common.keyword import validate_keyword_json 

11from sensorkit.core.client import SensorKit 

12 

13 

14@sk.declare_keyword 

15class BasicSafety(BaseModel): 

16 """Simple safety monitor keyword indicating whether conditions are currently safe to operate.""" 

17 is_safe: bool 

18 

19 

20SafetyProvider = sk.declare_trait( 

21 "SafetyProvider", 

22 required_keywords=("BasicSafety",), 

23) 

24 

25StandardSafety = sk.declare_archetype( 

26 "safety", 

27 required_traits=(SafetyProvider,), 

28) 

29"""Standard archetype for safety / go-no-go providers.""" 

30 

31 

32class SafetyConstraint(Constraint): 

33 """Constraint that monitors a BasicSafety provider and activates when conditions are unsafe.""" 

34 

35 kind: Literal["safety"] = "safety" 

36 provider: str 

37 

38 @override 

39 async def check_task(self, evaluator: ConstraintEvaluator, kit: SensorKit): 

40 provider = kit.entity(self.provider) 

41 consumer = await provider._stream.consume("BasicSafety") 

42 

43 async for msg in consumer: 

44 try: 

45 safety = validate_keyword_json("BasicSafety", msg.data) 

46 except Exception: 

47 continue 

48 

49 if not safety.is_safe: 

50 evaluator.constrain("unsafe") 

51 else: 

52 evaluator.clear("safe") 

53 

54 evaluator.ready()