Coverage for core / src / sensorkit / core / entity.py: 91%

1# SPDX-License-Identifier: Apache-2.0 

2from __future__ import annotations 

3 

4import asyncio 

5import contextlib 

6import functools 

7from abc import ABC, abstractmethod 

8from typing import TYPE_CHECKING, Annotated, Literal 

9 

10from loguru import logger 

11from pydantic import BaseModel, PlainSerializer, PlainValidator, ValidationError 

12 

13from sensorkit.backend.base import Entity 

14from sensorkit.backend.event import Event, EventMultiplexer, EventStreamConsumer 

15from sensorkit.backend.request import ( 

16 Call, 

17 ExtendedCall, 

18 ExtendedHandlerFunc, 

19 HandlerFunc, 

20 Request, 

21) 

22from sensorkit.common.aio import AsyncObserver 

23from sensorkit.common.keyword import ( 

24 Keyword, 

25 declare_keyword, 

26 get_keyword_info, 

27 validate_keyword_json, 

28) 

29from sensorkit.core.trait import match_archetype, match_traits 

30from sensorkit.data.graph import DataGraph 

31 

32if TYPE_CHECKING: 

33 from sensorkit.core.client import SensorKit 

34 

35type EntityType = Literal["generic", "device", "controller", "program"] 

36 

37 

38class DeviceDetails(BaseModel): 

39 """Device info.""" 

40 

41 supported_commands: frozenset[str] 

42 """Identifiers of each command supported by this device.""" 

43 

44 published_keywords: frozenset[str] 

45 """Identifiers of each keyword this device declares it publishes.""" 

46 

47 @functools.cached_property 

48 def archetype(self): 

49 return match_archetype(self) 

50 

51 @functools.cached_property 

52 def traits(self): 

53 return match_traits(self, exclude_archetypes=True) 

54 

55 

56class ControllerDetails(BaseModel): 

57 """Controller info.""" 

58 

59 supported_tasks: list[str] 

60 """Identifiers of each task supported by this controller.""" 

61 controlled_devices: list[str] 

62 """Names of each device commanded by this controller.""" 

63 

64 

65@declare_keyword 

66class EntityInfo(BaseModel): 

67 """Keyword describing an entity's type and capability details, stored in the KV backend.""" 

68 

69 entity_type: EntityType 

70 details: DeviceDetails | ControllerDetails | None 

71 

72 

73class EntityBase: 

74 """Base object that represents an entity exposed via a backend.""" 

75 

76 def __init__(self, sensorkit: SensorKit, entity: Entity): 

77 self._sensorkit = sensorkit 

78 self.backend = sensorkit.backend 

79 self.entity = entity 

80 

81 def sensorkit(self) -> SensorKit: 

82 """Return the SensorKit instance backing this entity.""" 

83 return self._sensorkit 

84 

85 @functools.cached_property 

86 def _request(self): 

87 return self.backend.request(self.entity) 

88 

89 @functools.cached_property 

90 def _stream(self): 

91 return self.backend.stream(self.entity) 

92 

93 @functools.cached_property 

94 def _kv(self): 

95 return self.backend.key_value(self.entity) 

96 

97 async def kv_put_model(self, model: BaseModel): 

98 """Serialise model to JSON and write it to the KV store under its class name.""" 

99 await self._kv.update( 

100 key=model.__class__.__name__, 

101 value=model.model_dump_json().encode(), 

102 ) 

103 

104 async def kv_get_model[M: BaseModel](self, model_type: type[M]) -> M: 

105 """Fetch and deserialise a model from the KV store by its class name.""" 

106 entry = await self._kv.get(model_type.__name__) 

107 return model_type.model_validate_json(entry.value) 

108 

109 async def kv_monitor_model[M: BaseModel](self, model_type: type[M]): 

110 """Monitor changes to a model in the KV store by its class name.""" 

111 stream = await self._kv.monitor(model_type.__name__) 

112 

113 async for entry in stream: 

114 if not entry.deleted(): 

115 with contextlib.suppress(ValidationError): 

116 yield model_type.model_validate_json(entry.value) 

117 

118 

119class EntityClient(EntityBase): 

120 """Object that exposes client-side functionality of an entity.""" 

121 

122 def __init__(self, sensorkit: SensorKit, entity: Entity): 

123 super().__init__(sensorkit, entity) 

124 self._online_monitor: asyncio.Task | None = None 

125 self._online_observer = AsyncObserver(False) 

126 

127 async def _monitor_online_state(self): 

128 stream = await self._kv.monitor("EntityLease") 

129 

130 async for entry in stream: 

131 online = not entry.deleted() 

132 

133 if self._online_observer.value != online: 

134 logger.debug(f"{self.entity} is now {'offline' if entry.deleted() else 'online'}") 

135 self._online_observer.notify(online) 

136 

137 def observe_online_state(self): 

138 """Return an async generator that yields True/False as the entity comes online or goes offline.""" 

139 if not self._online_monitor: 

140 self._online_monitor = asyncio.create_task(self._monitor_online_state()) 

141 

142 return self._online_observer.consume(initial_value=True) 

143 

144 @functools.cache 

145 def get_event_mux(self) -> EventMultiplexer: 

146 """Return the shared event consumer object for this entity.""" 

147 ec = EventStreamConsumer(self._stream) 

148 self._ec_startup = asyncio.create_task(ec.start()) 

149 return ec 

150 

151 @staticmethod 

152 async def _monitor_receive[M: Event](queue: asyncio.Queue[M], context: contextlib.ExitStack): 

153 with context: 

154 while True: 

155 yield await queue.get() 

156 queue.task_done() 

157 

158 async def _await_event_consumer(self, context: contextlib.ExitStack): 

159 """Wait until the entity's event subscription is live, releasing the queue on failure.""" 

160 try: 

161 await self._ec_startup 

162 except BaseException: 

163 context.close() 

164 raise 

165 

166 async def monitor_event[M: Event](self, event_type: type[M]): 

167 """Return an async generator yielding events of the specified type from the entity's stream.""" 

168 consumer = self.get_event_mux() 

169 context = contextlib.ExitStack() 

170 queue: asyncio.Queue[M] = context.enter_context(consumer.event_queue(event_type)) 

171 

172 await self._await_event_consumer(context) 

173 

174 return self._monitor_receive(queue, context) 

175 

176 async def monitor_all_events(self): 

177 """Return an async generator yielding all events from the entity's stream.""" 

178 consumer = self.get_event_mux() 

179 context = contextlib.ExitStack() 

180 queue = context.enter_context(consumer.all_events()) 

181 

182 await self._await_event_consumer(context) 

183 

184 return self._monitor_receive(queue, context) 

185 

186 async def monitor[M: Keyword | BaseModel](self, model_type: type[M]): 

187 """Return an async generator yielding (subject, model) tuples for the given keyword or model type.""" 

188 key = info.key if (info := get_keyword_info(model_type)) else None 

189 

190 if key is not None: 

191 validate_func = functools.partial(validate_keyword_json, key) 

192 else: 

193 validate_func = model_type.model_validate_json 

194 

195 consumer = await self._stream.consume(key, include_latest=True) 

196 

197 async def receive_data(): 

198 async for msg in consumer: 

199 try: 

200 model: M = validate_func(msg.data) 

201 yield msg.subject, model 

202 except ValidationError: 

203 logger.exception("Data validation failed") 

204 except Exception: 

205 logger.exception("Error while monitoring") 

206 

207 return receive_data() 

208 

209 async def monitor_all(self): 

210 """Return an async generator yielding (subject, model) tuples for all keyword updates on this entity.""" 

211 consumer = await self._stream.consume(include_latest=True) 

212 

213 async def receive_data(): 

214 async for msg in consumer: 

215 try: 

216 model = validate_keyword_json(msg.subject.prop, msg.data) 

217 except ValidationError: 

218 logger.exception("Data validation failed") 

219 continue 

220 

221 try: 

222 yield msg.subject, model 

223 except Exception: 

224 logger.exception("Error while monitoring") 

225 

226 return receive_data() 

227 

228 async def request[M: BaseModel](self, name: str, data: BaseModel, response_type: type[M]) -> M: 

229 """Send a named raw request and return the deserialised response.""" 

230 received = await self._request.invoke( 

231 name=name, 

232 payload=data.model_dump_json().encode(), 

233 ) 

234 return response_type.model_validate_json(received) 

235 

236 def call[P: BaseModel | None, R: BaseModel | None, V: BaseModel | None]( 

237 self, 

238 request: Request[P, R, V], 

239 data: P, 

240 ) -> Call[R, V]: 

241 """Invoke a request.""" 

242 payload = b"" if data is None else data.model_dump_json().encode() 

243 request_coro = self._request.invoke( 

244 name=request.name, 

245 payload=payload, 

246 ) 

247 

248 if request.is_extended(): 

249 return ExtendedCall( 

250 request_coro, 

251 request.response, 

252 request.result, 

253 self.get_event_mux(), 

254 ) 

255 else: 

256 return Call(request_coro, request.response) 

257 

258 

259class EntityRef[T: EntityClient = EntityClient]: 

260 """A serializable reference to an entity client.""" 

261 

262 name: str | None 

263 

264 def __init__(self, name: str | None = None): 

265 """Create a reference from an entity name. 

266 

267 Args: 

268 name: the serialized entity name, or `None` for an unset reference. 

269 """ 

270 self.name = name 

271 self._client: T | None = None 

272 

273 def _get_client(self, kit: SensorKit) -> T: 

274 """Return the client of the appropriate type for this reference.""" 

275 return kit.entity(self.name) 

276 

277 def resolve(self, kit: SensorKit) -> None: 

278 """Resolve this reference against a SensorKit instance. 

279 

280 If `name` is set, caches the corresponding entity client for later access via 

281 `get`, `require`, or `__call__`. 

282 """ 

283 if self.name is not None: 

284 self._client = self._get_client(kit) 

285 

286 def get(self) -> T | None: 

287 """Return the resolved client, or `None` if the reference is unset. 

288 

289 Raises: 

290 RuntimeError: if the reference has a name but has not yet been resolved. 

291 """ 

292 if self.name is not None and self._client is None: 

293 raise RuntimeError("Entity reference not resolved") 

294 return self._client 

295 

296 def require(self) -> T: 

297 """Return the resolved client, requiring that the reference be usable. 

298 

299 Raises: 

300 RuntimeError: if the reference is unset or has not been resolved. 

301 """ 

302 obj = self.get() 

303 if obj is None: 

304 raise RuntimeError("Required entity reference was not defined") 

305 return obj 

306 

307 def __call__(self) -> T | None: 

308 """Convenience accessor for `get`.""" 

309 return self.get() 

310 

311 def __eq__(self, other: object) -> bool: 

312 """Compare references by the entity they name.""" 

313 return isinstance(other, EntityRef) and self.name == other.name 

314 

315 def __hash__(self) -> int: 

316 return hash(self.name) 

317 

318 def __repr__(self) -> str: 

319 return f"{type(self).__name__}({self.name!r})" 

320 

321 @classmethod 

322 def __get_pydantic_core_schema__(cls, _source_type, handler): 

323 # Treat this reference as a string-like value. 

324 return handler( 

325 Annotated[ 

326 str, 

327 PlainValidator(lambda obj: cls(name=obj) if isinstance(obj, str) else obj), 

328 PlainSerializer(lambda obj: obj.name), 

329 ] 

330 ) 

331 

332 @classmethod 

333 def __get_pydantic_json_schema__(cls, core_schema, handler): 

334 schema = handler(core_schema) 

335 schema["type"] = "string" 

336 return schema 

337 

338 

339class EntityInterface(ABC): 

340 """Interface describing an implementation of an entity.""" 

341 

342 @abstractmethod 

343 def sensorkit(self) -> SensorKit: 

344 """Return the SensorKit instance backing this entity.""" 

345 ... 

346 

347 @property 

348 @abstractmethod 

349 def task_group(self) -> asyncio.TaskGroup: 

350 """Return the TaskGroup used for background tasks on this entity.""" 

351 ... 

352 

353 @abstractmethod 

354 async def kv_put_model(self, model: BaseModel): 

355 """Serialise model and write it to the KV store.""" 

356 ... 

357 

358 @abstractmethod 

359 async def kv_get_model[M: BaseModel](self, model_type: type[M]) -> M: 

360 """Fetch and deserialise a model from the KV store.""" 

361 ... 

362 

363 @abstractmethod 

364 async def emit_event(self, event: Event): 

365 """Publish an event to the entity's event stream.""" 

366 ... 

367 

368 @abstractmethod 

369 async def publish(self, model: Keyword): 

370 """Publish a keyword model to the entity's data stream.""" 

371 ... 

372 

373 @abstractmethod 

374 async def handle_request[P: BaseModel | None, R: BaseModel | None, V: BaseModel | None]( 

375 self, 

376 request: Request[P, R, V], 

377 func: HandlerFunc[P, R] | ExtendedHandlerFunc[P, R, V], 

378 ): 

379 """Register a handler for the given Request definition.""" 

380 ... 

381 

382 @abstractmethod 

383 async def data_graph(self) -> DataGraph: 

384 """Return the DataGraph for this entity, creating it if necessary.""" 

385 ... 

386 

387 @abstractmethod 

388 async def publish_entity_info(self) -> EntityInfo: 

389 """Build and publish EntityInfo to the KV store.""" 

390 ...