digitalkin 0.3.1.dev2__py3-none-any.whl → 0.3.2a3__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (87) hide show
  1. base_server/server_async_insecure.py +6 -5
  2. base_server/server_async_secure.py +6 -5
  3. base_server/server_sync_insecure.py +5 -4
  4. base_server/server_sync_secure.py +5 -4
  5. digitalkin/__version__.py +1 -1
  6. digitalkin/core/job_manager/base_job_manager.py +1 -1
  7. digitalkin/core/job_manager/single_job_manager.py +78 -36
  8. digitalkin/core/job_manager/taskiq_broker.py +7 -6
  9. digitalkin/core/job_manager/taskiq_job_manager.py +9 -5
  10. digitalkin/core/task_manager/base_task_manager.py +3 -1
  11. digitalkin/core/task_manager/surrealdb_repository.py +29 -7
  12. digitalkin/core/task_manager/task_executor.py +46 -12
  13. digitalkin/core/task_manager/task_session.py +132 -102
  14. digitalkin/grpc_servers/module_server.py +95 -171
  15. digitalkin/grpc_servers/module_servicer.py +121 -19
  16. digitalkin/grpc_servers/utils/grpc_client_wrapper.py +36 -10
  17. digitalkin/grpc_servers/utils/utility_schema_extender.py +106 -0
  18. digitalkin/models/__init__.py +1 -1
  19. digitalkin/models/core/job_manager_models.py +0 -8
  20. digitalkin/models/core/task_monitor.py +23 -1
  21. digitalkin/models/grpc_servers/models.py +95 -8
  22. digitalkin/models/module/__init__.py +26 -13
  23. digitalkin/models/module/base_types.py +61 -0
  24. digitalkin/models/module/module_context.py +279 -13
  25. digitalkin/models/module/module_types.py +28 -392
  26. digitalkin/models/module/setup_types.py +547 -0
  27. digitalkin/models/module/tool_cache.py +230 -0
  28. digitalkin/models/module/tool_reference.py +160 -0
  29. digitalkin/models/module/utility.py +167 -0
  30. digitalkin/models/services/cost.py +22 -1
  31. digitalkin/models/services/registry.py +77 -0
  32. digitalkin/modules/__init__.py +5 -1
  33. digitalkin/modules/_base_module.py +188 -63
  34. digitalkin/modules/archetype_module.py +6 -1
  35. digitalkin/modules/tool_module.py +6 -1
  36. digitalkin/modules/triggers/__init__.py +8 -0
  37. digitalkin/modules/triggers/healthcheck_ping_trigger.py +45 -0
  38. digitalkin/modules/triggers/healthcheck_services_trigger.py +63 -0
  39. digitalkin/modules/triggers/healthcheck_status_trigger.py +52 -0
  40. digitalkin/services/__init__.py +4 -0
  41. digitalkin/services/communication/__init__.py +7 -0
  42. digitalkin/services/communication/communication_strategy.py +87 -0
  43. digitalkin/services/communication/default_communication.py +104 -0
  44. digitalkin/services/communication/grpc_communication.py +264 -0
  45. digitalkin/services/cost/cost_strategy.py +36 -14
  46. digitalkin/services/cost/default_cost.py +61 -1
  47. digitalkin/services/cost/grpc_cost.py +98 -2
  48. digitalkin/services/filesystem/grpc_filesystem.py +9 -2
  49. digitalkin/services/registry/__init__.py +22 -1
  50. digitalkin/services/registry/default_registry.py +156 -4
  51. digitalkin/services/registry/exceptions.py +47 -0
  52. digitalkin/services/registry/grpc_registry.py +382 -0
  53. digitalkin/services/registry/registry_models.py +15 -0
  54. digitalkin/services/registry/registry_strategy.py +106 -4
  55. digitalkin/services/services_config.py +25 -3
  56. digitalkin/services/services_models.py +5 -1
  57. digitalkin/services/setup/default_setup.py +1 -1
  58. digitalkin/services/setup/grpc_setup.py +1 -1
  59. digitalkin/services/storage/grpc_storage.py +1 -1
  60. digitalkin/services/user_profile/__init__.py +11 -0
  61. digitalkin/services/user_profile/grpc_user_profile.py +2 -2
  62. digitalkin/services/user_profile/user_profile_strategy.py +0 -15
  63. digitalkin/utils/__init__.py +15 -3
  64. digitalkin/utils/conditional_schema.py +260 -0
  65. digitalkin/utils/dynamic_schema.py +4 -0
  66. digitalkin/utils/schema_splitter.py +290 -0
  67. {digitalkin-0.3.1.dev2.dist-info → digitalkin-0.3.2a3.dist-info}/METADATA +12 -12
  68. digitalkin-0.3.2a3.dist-info/RECORD +144 -0
  69. {digitalkin-0.3.1.dev2.dist-info → digitalkin-0.3.2a3.dist-info}/WHEEL +1 -1
  70. {digitalkin-0.3.1.dev2.dist-info → digitalkin-0.3.2a3.dist-info}/top_level.txt +1 -0
  71. modules/archetype_with_tools_module.py +232 -0
  72. modules/cpu_intensive_module.py +1 -1
  73. modules/dynamic_setup_module.py +5 -29
  74. modules/minimal_llm_module.py +1 -1
  75. modules/text_transform_module.py +1 -1
  76. monitoring/digitalkin_observability/__init__.py +46 -0
  77. monitoring/digitalkin_observability/http_server.py +150 -0
  78. monitoring/digitalkin_observability/interceptors.py +176 -0
  79. monitoring/digitalkin_observability/metrics.py +201 -0
  80. monitoring/digitalkin_observability/prometheus.py +137 -0
  81. monitoring/tests/test_metrics.py +172 -0
  82. services/filesystem_module.py +7 -5
  83. services/storage_module.py +4 -2
  84. digitalkin/grpc_servers/registry_server.py +0 -65
  85. digitalkin/grpc_servers/registry_servicer.py +0 -456
  86. digitalkin-0.3.1.dev2.dist-info/RECORD +0 -119
  87. {digitalkin-0.3.1.dev2.dist-info → digitalkin-0.3.2a3.dist-info}/licenses/LICENSE +0 -0
@@ -2,24 +2,28 @@
2
2
 
3
3
  import asyncio
4
4
  import json
5
+ import os
5
6
  from abc import ABC, abstractmethod
6
7
  from collections.abc import Callable, Coroutine
7
8
  from typing import Any, ClassVar, Generic
8
9
 
10
+ from digitalkin.grpc_servers.utils.utility_schema_extender import UtilitySchemaExtender
9
11
  from digitalkin.logger import logger
10
- from digitalkin.models.module import (
12
+ from digitalkin.models.module.module import ModuleCodeModel, ModuleStatus
13
+ from digitalkin.models.module.module_context import ModuleContext
14
+ from digitalkin.models.module.module_types import (
15
+ DataModel,
11
16
  InputModelT,
12
- ModuleStatus,
13
17
  OutputModelT,
14
18
  SecretModelT,
15
19
  SetupModelT,
16
20
  )
17
- from digitalkin.models.module.module import ModuleCodeModel
18
- from digitalkin.models.module.module_context import ModuleContext
21
+ from digitalkin.models.module.utility import EndOfStreamOutput, ModuleStartInfoOutput, UtilityProtocol
22
+ from digitalkin.models.services.storage import BaseRole
19
23
  from digitalkin.modules.trigger_handler import TriggerHandler
20
24
  from digitalkin.services.services_config import ServicesConfig, ServicesStrategy
21
- from digitalkin.utils.llm_ready_schema import llm_ready_schema
22
25
  from digitalkin.utils.package_discover import ModuleDiscoverer
26
+ from digitalkin.utils.schema_splitter import SchemaSplitter
23
27
 
24
28
 
25
29
  class BaseModule( # noqa: PLR0904
@@ -46,10 +50,20 @@ class BaseModule( # noqa: PLR0904
46
50
  triggers_discoverer: ClassVar[ModuleDiscoverer]
47
51
 
48
52
  # service config params
49
- services_config_strategies: ClassVar[dict[str, ServicesStrategy | None]]
50
- services_config_params: ClassVar[dict[str, dict[str, Any | None] | None]]
53
+ services_config_strategies: ClassVar[dict[str, ServicesStrategy | None]] = {}
54
+ services_config_params: ClassVar[dict[str, dict[str, Any | None] | None]] = {}
51
55
  services_config: ServicesConfig
52
56
 
57
+ @classmethod
58
+ def get_module_id(cls) -> str:
59
+ """Get the module ID from environment variable or metadata.
60
+
61
+ Returns:
62
+ The module_id from DIGITALKIN_MODULE_ID env var, or metadata module_id,
63
+ or "unknown" if neither exists.
64
+ """
65
+ return os.environ.get("DIGITALKIN_MODULE_ID") or cls.metadata.get("module_id", "unknown")
66
+
53
67
  def _init_strategies(self, mission_id: str, setup_id: str, setup_version_id: str) -> dict[str, Any]:
54
68
  """Initialize the services configuration.
55
69
 
@@ -62,6 +76,7 @@ class BaseModule( # noqa: PLR0904
62
76
  registry: RegistryStrategy
63
77
  snapshot: SnapshotStrategy
64
78
  storage: StorageStrategy
79
+ user_profile: UserProfileStrategy
65
80
  """
66
81
  logger.debug("Service initialisation: %s", self.services_config_strategies.keys())
67
82
  return {
@@ -122,7 +137,8 @@ class BaseModule( # noqa: PLR0904
122
137
  """
123
138
  if cls.secret_format is not None:
124
139
  if llm_format:
125
- return json.dumps(llm_ready_schema(cls.secret_format), indent=2)
140
+ result_json, result_ui = SchemaSplitter.split(cls.secret_format.model_json_schema())
141
+ return json.dumps({"json_schema": result_json, "ui_schema": result_ui}, indent=2)
126
142
  return json.dumps(cls.secret_format.model_json_schema(), indent=2)
127
143
  msg = f"{cls.__name__}' class does not define a 'secret_format'."
128
144
  raise NotImplementedError(msg)
@@ -141,12 +157,16 @@ class BaseModule( # noqa: PLR0904
141
157
  Raises:
142
158
  NotImplementedError: If the `input_format` class attribute is not defined.
143
159
  """
144
- if cls.input_format is not None:
145
- if llm_format:
146
- return json.dumps(llm_ready_schema(cls.input_format), indent=2)
147
- return json.dumps(cls.input_format.model_json_schema(), indent=2)
148
- msg = f"{cls.__name__}' class does not define an 'input_format'."
149
- raise NotImplementedError(msg)
160
+ if cls.input_format is None:
161
+ msg = f"{cls.__name__}' class does not define an 'input_format'."
162
+ raise NotImplementedError(msg)
163
+
164
+ extended_model = UtilitySchemaExtender.create_extended_input_model(cls.input_format)
165
+
166
+ if llm_format:
167
+ result_json, result_ui = SchemaSplitter.split(extended_model.model_json_schema())
168
+ return json.dumps({"json_schema": result_json, "ui_schema": result_ui}, indent=2)
169
+ return json.dumps(extended_model.model_json_schema(), indent=2)
150
170
 
151
171
  @classmethod
152
172
  async def get_output_format(cls, *, llm_format: bool) -> str:
@@ -162,12 +182,16 @@ class BaseModule( # noqa: PLR0904
162
182
  Raises:
163
183
  NotImplementedError: If the `output_format` class attribute is not defined.
164
184
  """
165
- if cls.output_format is not None:
166
- if llm_format:
167
- return json.dumps(llm_ready_schema(cls.output_format), indent=2)
168
- return json.dumps(cls.output_format.model_json_schema(), indent=2)
169
- msg = "'%s' class does not define an 'output_format'."
170
- raise NotImplementedError(msg)
185
+ if cls.output_format is None:
186
+ msg = f"'{cls.__name__}' class does not define an 'output_format'."
187
+ raise NotImplementedError(msg)
188
+
189
+ extended_model = UtilitySchemaExtender.create_extended_output_model(cls.output_format)
190
+
191
+ if llm_format:
192
+ result_json, result_ui = SchemaSplitter.split(extended_model.model_json_schema())
193
+ return json.dumps({"json_schema": result_json, "ui_schema": result_ui}, indent=2)
194
+ return json.dumps(extended_model.model_json_schema(), indent=2)
171
195
 
172
196
  @classmethod
173
197
  async def get_config_setup_format(cls, *, llm_format: bool) -> str:
@@ -194,7 +218,8 @@ class BaseModule( # noqa: PLR0904
194
218
  if cls.setup_format is not None:
195
219
  setup_format = await cls.setup_format.get_clean_model(config_fields=True, hidden_fields=False, force=True)
196
220
  if llm_format:
197
- return json.dumps(llm_ready_schema(setup_format), indent=2)
221
+ result_json, result_ui = SchemaSplitter.split(setup_format.model_json_schema())
222
+ return json.dumps({"json_schema": result_json, "ui_schema": result_ui}, indent=2)
198
223
  return json.dumps(setup_format.model_json_schema(), indent=2)
199
224
  msg = "'%s' class does not define an 'config_setup_format'."
200
225
  raise NotImplementedError(msg)
@@ -223,11 +248,49 @@ class BaseModule( # noqa: PLR0904
223
248
  if cls.setup_format is not None:
224
249
  setup_format = await cls.setup_format.get_clean_model(config_fields=False, hidden_fields=True, force=True)
225
250
  if llm_format:
226
- return json.dumps(llm_ready_schema(setup_format), indent=2)
251
+ result_json, result_ui = SchemaSplitter.split(setup_format.model_json_schema())
252
+ return json.dumps({"json_schema": result_json, "ui_schema": result_ui}, indent=2)
227
253
  return json.dumps(setup_format.model_json_schema(), indent=2)
228
254
  msg = "'%s' class does not define an 'setup_format'."
229
255
  raise NotImplementedError(msg)
230
256
 
257
+ @classmethod
258
+ async def get_cost_format(cls, *, llm_format: bool) -> str:
259
+ """Get the JSON schema of the cost configuration.
260
+
261
+ Extracts CostConfig from services_config_params["cost"]["config"]
262
+ and returns as JSON schema.
263
+
264
+ Args:
265
+ llm_format: If True, return LLM-optimized schema format with inlined
266
+ references and simplified structure.
267
+
268
+ Returns:
269
+ The JSON schema of the cost configuration as a JSON string.
270
+ """
271
+ cost_params = cls.services_config_params.get("cost", {})
272
+ config = cost_params.get("config", {}) if cost_params else {}
273
+
274
+ if not config:
275
+ return json.dumps({}, indent=2)
276
+
277
+ # Convert CostConfig objects to serializable dict
278
+ cost_schema = {
279
+ name: {
280
+ "name": cost_config.name,
281
+ "type": cost_config.type.value if hasattr(cost_config.type, "value") else cost_config.type,
282
+ "description": cost_config.description,
283
+ "unit": cost_config.unit,
284
+ "rate": cost_config.rate,
285
+ }
286
+ for name, cost_config in config.items()
287
+ }
288
+
289
+ if llm_format:
290
+ result_json, result_ui = SchemaSplitter.split({"costs": cost_schema})
291
+ return json.dumps({"json_schema": result_json, "ui_schema": result_ui}, indent=2)
292
+ return json.dumps(cost_schema, indent=2)
293
+
231
294
  @classmethod
232
295
  def create_config_setup_model(cls, config_setup_data: dict[str, Any]) -> SetupModelT:
233
296
  """Create the setup model from the setup data.
@@ -304,8 +367,18 @@ class BaseModule( # noqa: PLR0904
304
367
  If a package is provided, all .py files within its path are imported; otherwise, the current
305
368
  working directory is searched. For each imported module, any class matching the criteria is
306
369
  registered via cls.register(). Errors during import are logged at debug level.
370
+
371
+ Built-in healthcheck handlers (ping, services, status) are automatically registered
372
+ to provide standard healthcheck functionality for all modules.
307
373
  """
374
+ from digitalkin.models.module.utility import UtilityRegistry # noqa: PLC0415
375
+
308
376
  cls.triggers_discoverer.discover_modules()
377
+
378
+ # Auto-register built-in SDK triggers (healthcheck, etc.)
379
+ for trigger_cls in UtilityRegistry.get_builtin_triggers():
380
+ cls.triggers_discoverer.register_trigger(trigger_cls)
381
+
309
382
  logger.debug("discovered: %s", cls.triggers_discoverer)
310
383
 
311
384
  @classmethod
@@ -320,25 +393,6 @@ class BaseModule( # noqa: PLR0904
320
393
  """
321
394
  return cls.triggers_discoverer.register_trigger(handler_cls)
322
395
 
323
- async def run_config_setup( # noqa: PLR6301
324
- self,
325
- context: ModuleContext, # noqa: ARG002
326
- config_setup_data: SetupModelT,
327
- ) -> SetupModelT:
328
- """Run config setup the module.
329
-
330
- The config setup is used to initialize the setup with configuration data.
331
- This method is typically used to set up the module with necessary configuration before running it,
332
- especially for processing data like files.
333
- The function needs to save the setup in the storage.
334
- The module will be initialize with the setup and not the config setup.
335
- This method is optional, the config setup and setup can be the same.
336
-
337
- Returns:
338
- The updated setup model after running the config setup.
339
- """
340
- return config_setup_data
341
-
342
396
  @abstractmethod
343
397
  async def initialize(self, context: ModuleContext, setup_data: SetupModelT) -> None:
344
398
  """Initialize the module."""
@@ -349,24 +403,22 @@ class BaseModule( # noqa: PLR0904
349
403
  input_data: InputModelT,
350
404
  setup_data: SetupModelT,
351
405
  ) -> None:
352
- """Run the module with the given input and setup data.
353
-
354
- This method validates the input data, determines the protocol from the input,
355
- and dispatches the request to the corresponding trigger handler. The trigger handler
356
- is responsible for processing the input and invoking the callback with the result.
357
-
358
- Triggers:
359
- - The method is triggered when a module run is requested with specific input and setup data.
360
- - The protocol specified in the input determines which trigger handler is invoked.
406
+ """Run the module by dispatching to the appropriate trigger handler.
361
407
 
362
408
  Args:
363
- input_data (InputModelT): The input data to be processed by the module.
364
- setup_data (SetupModelT): The setup or configuration data required for the module.
409
+ input_data: Input data to process.
410
+ setup_data: Configuration data for the module.
365
411
 
366
412
  Raises:
367
413
  ValueError: If no handler for the protocol is found.
368
414
  """
369
415
  input_instance = self.input_format.model_validate(input_data)
416
+
417
+ # Apply cost limits if present in input
418
+ cost_limits = getattr(input_instance, "cost_limits", None)
419
+ if cost_limits is not None and self.context.cost is not None:
420
+ self.context.cost.set_limits(cost_limits)
421
+
370
422
  handler_instance = self.triggers_discoverer.get_trigger(
371
423
  input_instance.root.protocol,
372
424
  input_instance.root,
@@ -383,6 +435,25 @@ class BaseModule( # noqa: PLR0904
383
435
  """Run the module."""
384
436
  raise NotImplementedError
385
437
 
438
+ async def run_config_setup( # noqa: PLR6301
439
+ self,
440
+ context: ModuleContext, # noqa: ARG002
441
+ config_setup_data: SetupModelT,
442
+ ) -> SetupModelT:
443
+ """Run config setup the module.
444
+
445
+ The config setup is used to initialize the setup with configuration data.
446
+ This method is typically used to set up the module with necessary configuration before running it,
447
+ especially for processing data like files.
448
+ The function needs to save the setup in the storage.
449
+ The module will be initialize with the setup and not the config setup.
450
+ This method is optional, the config setup and setup can be the same.
451
+
452
+ Returns:
453
+ The updated setup model after running the config setup.
454
+ """
455
+ return config_setup_data
456
+
386
457
  async def _run_lifecycle(
387
458
  self,
388
459
  input_data: InputModelT,
@@ -410,13 +481,32 @@ class BaseModule( # noqa: PLR0904
410
481
  self,
411
482
  input_data: InputModelT,
412
483
  setup_data: SetupModelT,
413
- callback: Callable[[OutputModelT | ModuleCodeModel], Coroutine[Any, Any, None]],
484
+ callback: Callable[[OutputModelT | ModuleCodeModel | DataModel[UtilityProtocol]], Coroutine[Any, Any, None]],
414
485
  done_callback: Callable | None = None,
415
486
  ) -> None:
416
487
  """Start the module."""
417
488
  try:
418
489
  self.context.callbacks.send_message = callback
419
- logger.info(f"Inititalize module {self.context.session.job_id}")
490
+
491
+ tool_cache = setup_data.build_tool_cache()
492
+ if tool_cache.entries:
493
+ self.context.tool_cache = tool_cache
494
+
495
+ await callback(
496
+ DataModel(
497
+ root=ModuleStartInfoOutput(
498
+ job_id=self.context.session.job_id,
499
+ mission_id=self.context.session.mission_id,
500
+ setup_id=self.context.session.setup_id,
501
+ setup_version_id=self.context.session.setup_version_id,
502
+ module_id=self.get_module_id(),
503
+ module_name=self.name,
504
+ ),
505
+ annotations={"role": BaseRole.SYSTEM},
506
+ )
507
+ )
508
+
509
+ logger.info("Initialize module %s", self.context.session.job_id)
420
510
  await self.initialize(self.context, setup_data)
421
511
  except Exception as e:
422
512
  self._status = ModuleStatus.FAILED
@@ -437,7 +527,7 @@ class BaseModule( # noqa: PLR0904
437
527
  try:
438
528
  logger.debug("Init the discovered input handlers.")
439
529
  self.triggers_discoverer.init_handlers(self.context)
440
- logger.debug(f"Run lifecycle {self.context.session.job_id}")
530
+ logger.debug("Run lifecycle %s", self.context.session.job_id)
441
531
  await self._run_lifecycle(input_data, setup_data)
442
532
  except Exception:
443
533
  self._status = ModuleStatus.FAILED
@@ -452,31 +542,66 @@ class BaseModule( # noqa: PLR0904
452
542
  self._status = ModuleStatus.STOPPING
453
543
  logger.debug("Module %s stopped", self.name)
454
544
  await self.cleanup()
455
- await self.context.callbacks.send_message(ModuleCodeModel(code="__END_OF_STREAM__"))
545
+ await self.context.callbacks.send_message(
546
+ DataModel(
547
+ root=EndOfStreamOutput(),
548
+ annotations={"role": BaseRole.SYSTEM},
549
+ )
550
+ )
456
551
  self._status = ModuleStatus.STOPPED
457
552
  logger.debug("Module %s cleaned", self.name)
458
553
  except Exception:
459
554
  self._status = ModuleStatus.FAILED
460
555
  logger.exception("Error stopping module")
461
556
 
557
+ async def _resolve_tools(self, config_setup_data: SetupModelT) -> None:
558
+ """Resolve tool references and build cache.
559
+
560
+ Args:
561
+ config_setup_data: Setup data containing tool references.
562
+ """
563
+ logger.info("Starting tool resolution", extra=self.context.session.current_ids())
564
+ if self.context.registry is not None and self.context.communication is not None:
565
+ await config_setup_data.resolve_tool_references(self.context.registry, self.context.communication)
566
+ logger.info("Tool references resolved", extra=self.context.session.current_ids())
567
+ else:
568
+ logger.warning(
569
+ "No registry or communication available, skipping tool resolution",
570
+ extra=self.context.session.current_ids(),
571
+ )
572
+
573
+ tool_cache = config_setup_data.build_tool_cache()
574
+ self.context.tool_cache = tool_cache
575
+ logger.info(
576
+ "Tool cache built with %d entries: %s",
577
+ len(tool_cache.entries),
578
+ list(tool_cache.entries.keys()),
579
+ extra=self.context.session.current_ids(),
580
+ )
581
+
462
582
  async def start_config_setup(
463
583
  self,
464
584
  config_setup_data: SetupModelT,
465
585
  callback: Callable[[SetupModelT | ModuleCodeModel], Coroutine[Any, Any, None]],
466
586
  ) -> None:
467
- """Start the module."""
587
+ """Run config setup lifecycle with tool resolution in parallel.
588
+
589
+ Args:
590
+ config_setup_data: Initial setup data to configure.
591
+ callback: Callback to send the configured setup model.
592
+ """
468
593
  try:
469
594
  logger.info("Run Config Setup lifecycle", extra=self.context.session.current_ids())
470
595
  self._status = ModuleStatus.RUNNING
471
596
  self.context.callbacks.set_config_setup = callback
472
- content = await self.run_config_setup(self.context, config_setup_data)
473
597
 
474
- wrapper = config_setup_data.model_dump()
475
- wrapper["content"] = content.model_dump()
476
- setup_model = await self.create_setup_model(wrapper)
598
+ # Resolve tools first to populate companion fields, then run config setup
599
+ await self._resolve_tools(config_setup_data)
600
+ updated_config = await self.run_config_setup(self.context, config_setup_data)
601
+
602
+ setup_model = await self.create_setup_model(updated_config.model_dump())
477
603
  await callback(setup_model)
478
604
  self._status = ModuleStatus.STOPPING
479
605
  except Exception:
480
- logger.error("Error during module lifecyle")
481
606
  self._status = ModuleStatus.FAILED
482
- logger.exception("Error during module lifecyle", extra=self.context.session.current_ids())
607
+ logger.exception("Error during config setup lifecycle", extra=self.context.session.current_ids())
@@ -2,7 +2,12 @@
2
2
 
3
3
  from abc import ABC
4
4
 
5
- from digitalkin.models.module import InputModelT, OutputModelT, SecretModelT, SetupModelT
5
+ from digitalkin.models.module.module_types import (
6
+ InputModelT,
7
+ OutputModelT,
8
+ SecretModelT,
9
+ SetupModelT,
10
+ )
6
11
  from digitalkin.modules._base_module import BaseModule
7
12
 
8
13
 
@@ -2,7 +2,12 @@
2
2
 
3
3
  from abc import ABC
4
4
 
5
- from digitalkin.models.module import InputModelT, OutputModelT, SecretModelT, SetupModelT
5
+ from digitalkin.models.module.module_types import (
6
+ InputModelT,
7
+ OutputModelT,
8
+ SecretModelT,
9
+ SetupModelT,
10
+ )
6
11
  from digitalkin.modules._base_module import BaseModule # type: ignore
7
12
 
8
13
 
@@ -0,0 +1,8 @@
1
+ """Built-in SDK triggers.
2
+
3
+ These triggers are automatically registered without requiring discovery.
4
+ They provide standard functionality available to all modules.
5
+
6
+ Note: These are internal triggers. External code should not import them directly.
7
+ Use UtilityRegistry.get_builtin_triggers() to access the trigger classes.
8
+ """
@@ -0,0 +1,45 @@
1
+ """Healthcheck ping trigger - simple alive check."""
2
+
3
+ from datetime import datetime
4
+ from typing import Any, ClassVar
5
+
6
+ from digitalkin.mixins import BaseMixin
7
+ from digitalkin.models.module.module_context import ModuleContext
8
+ from digitalkin.models.module.utility import (
9
+ HealthcheckPingInput,
10
+ HealthcheckPingOutput,
11
+ )
12
+ from digitalkin.modules.trigger_handler import TriggerHandler
13
+
14
+
15
+ class HealthcheckPingTrigger(TriggerHandler, BaseMixin):
16
+ """Handler for simple ping healthcheck.
17
+
18
+ Responds immediately with "pong" status to verify the module is responsive.
19
+ """
20
+
21
+ protocol: ClassVar[str] = "healthcheck_ping"
22
+ input_format = HealthcheckPingInput
23
+ _request_time: datetime
24
+
25
+ def __init__(self, context: ModuleContext) -> None:
26
+ """Initialize the handler."""
27
+ self._request_time = datetime.now(tz=context.session.timezone)
28
+
29
+ async def handle(
30
+ self,
31
+ input_data: HealthcheckPingInput, # noqa: ARG002
32
+ setup_data: Any, # noqa: ANN401, ARG002
33
+ context: ModuleContext,
34
+ ) -> None:
35
+ """Handle ping healthcheck request.
36
+
37
+ Args:
38
+ input_data: The input trigger data (unused for healthcheck).
39
+ setup_data: The setup configuration (unused for healthcheck).
40
+ context: The module context.
41
+ """
42
+ elapsed = datetime.now(tz=context.session.timezone) - self._request_time
43
+ latency_ms = elapsed.total_seconds() * 1000
44
+ output = HealthcheckPingOutput(latency_ms=latency_ms)
45
+ await self.send_message(context, output)
@@ -0,0 +1,63 @@
1
+ """Healthcheck services trigger - reports service health."""
2
+
3
+ from typing import Any, ClassVar
4
+
5
+ from digitalkin.mixins import BaseMixin
6
+ from digitalkin.models.module.module_context import ModuleContext
7
+ from digitalkin.models.module.utility import (
8
+ HealthcheckServicesInput,
9
+ HealthcheckServicesOutput,
10
+ ServiceHealthStatus,
11
+ )
12
+ from digitalkin.modules.trigger_handler import TriggerHandler
13
+
14
+
15
+ class HealthcheckServicesTrigger(TriggerHandler, BaseMixin):
16
+ """Handler for services healthcheck.
17
+
18
+ Reports the health status of all configured services (storage, cost, filesystem, etc.).
19
+ """
20
+
21
+ protocol: ClassVar[str] = "healthcheck_services"
22
+ input_format = HealthcheckServicesInput
23
+
24
+ def __init__(self, context: ModuleContext) -> None:
25
+ """Initialize the handler."""
26
+
27
+ async def handle(
28
+ self,
29
+ input_data: HealthcheckServicesInput, # noqa: ARG002
30
+ setup_data: Any, # noqa: ANN401, ARG002
31
+ context: ModuleContext,
32
+ ) -> None:
33
+ """Handle services healthcheck request.
34
+
35
+ Args:
36
+ input_data: The input trigger data (unused for healthcheck).
37
+ setup_data: The setup configuration (unused for healthcheck).
38
+ context: The module context.
39
+ """
40
+ service_names = ["storage", "cost", "filesystem", "registry", "user_profile"]
41
+ services_status: list[ServiceHealthStatus] = []
42
+
43
+ for name in service_names:
44
+ service = getattr(context, name, None)
45
+ if service is not None:
46
+ services_status.append(ServiceHealthStatus(name=name, status="healthy"))
47
+ else:
48
+ services_status.append(ServiceHealthStatus(name=name, status="unknown", message="Not configured"))
49
+
50
+ # Determine overall status
51
+ healthy_count = sum(1 for s in services_status if s.status == "healthy")
52
+ if healthy_count == len(services_status):
53
+ overall_status = "healthy"
54
+ elif healthy_count > 0:
55
+ overall_status = "degraded"
56
+ else:
57
+ overall_status = "unhealthy"
58
+
59
+ output = HealthcheckServicesOutput(
60
+ services=services_status,
61
+ overall_status=overall_status, # type: ignore[arg-type]
62
+ )
63
+ await self.send_message(context, output)
@@ -0,0 +1,52 @@
1
+ """Healthcheck status trigger - comprehensive module status."""
2
+
3
+ import time
4
+ from typing import Any, ClassVar
5
+
6
+ from digitalkin.mixins import BaseMixin
7
+ from digitalkin.models.module.module_context import ModuleContext
8
+ from digitalkin.models.module.utility import (
9
+ HealthcheckStatusInput,
10
+ HealthcheckStatusOutput,
11
+ )
12
+ from digitalkin.modules.trigger_handler import TriggerHandler
13
+
14
+
15
+ class HealthcheckStatusTrigger(TriggerHandler, BaseMixin):
16
+ """Handler for comprehensive status healthcheck.
17
+
18
+ Reports detailed module status including uptime, active jobs, and metadata.
19
+ """
20
+
21
+ protocol: ClassVar[str] = "healthcheck_status"
22
+ input_format = HealthcheckStatusInput
23
+ _start_time: ClassVar[float] = time.time()
24
+
25
+ def __init__(self, context: ModuleContext) -> None:
26
+ """Initialize the handler."""
27
+
28
+ async def handle(
29
+ self,
30
+ input_data: HealthcheckStatusInput, # noqa: ARG002
31
+ setup_data: Any, # noqa: ANN401, ARG002
32
+ context: ModuleContext,
33
+ ) -> None:
34
+ """Handle status healthcheck request.
35
+
36
+ Args:
37
+ input_data: The input trigger data (unused for healthcheck).
38
+ setup_data: The setup configuration (unused for healthcheck).
39
+ context: The module context.
40
+ """
41
+ output = HealthcheckStatusOutput(
42
+ module_name=context.session.setup_id,
43
+ module_status="RUNNING",
44
+ uptime_seconds=time.time() - self._start_time,
45
+ active_jobs=1,
46
+ metadata={
47
+ "job_id": context.session.job_id,
48
+ "mission_id": context.session.mission_id,
49
+ "setup_version_id": context.session.setup_version_id,
50
+ },
51
+ )
52
+ await self.send_message(context, output)
@@ -1,6 +1,7 @@
1
1
  """This package contains the abstract base class for all services."""
2
2
 
3
3
  from digitalkin.services.agent import AgentStrategy, DefaultAgent
4
+ from digitalkin.services.communication import CommunicationStrategy, DefaultCommunication, GrpcCommunication
4
5
  from digitalkin.services.cost import CostStrategy, DefaultCost
5
6
  from digitalkin.services.filesystem import DefaultFilesystem, FilesystemStrategy
6
7
  from digitalkin.services.identity import DefaultIdentity, IdentityStrategy
@@ -10,8 +11,10 @@ from digitalkin.services.storage import DefaultStorage, StorageStrategy
10
11
 
11
12
  __all__ = [
12
13
  "AgentStrategy",
14
+ "CommunicationStrategy",
13
15
  "CostStrategy",
14
16
  "DefaultAgent",
17
+ "DefaultCommunication",
15
18
  "DefaultCost",
16
19
  "DefaultFilesystem",
17
20
  "DefaultIdentity",
@@ -19,6 +22,7 @@ __all__ = [
19
22
  "DefaultSnapshot",
20
23
  "DefaultStorage",
21
24
  "FilesystemStrategy",
25
+ "GrpcCommunication",
22
26
  "IdentityStrategy",
23
27
  "RegistryStrategy",
24
28
  "SnapshotStrategy",
@@ -0,0 +1,7 @@
1
+ """Communication service for module-to-module interaction."""
2
+
3
+ from digitalkin.services.communication.communication_strategy import CommunicationStrategy
4
+ from digitalkin.services.communication.default_communication import DefaultCommunication
5
+ from digitalkin.services.communication.grpc_communication import GrpcCommunication
6
+
7
+ __all__ = ["CommunicationStrategy", "DefaultCommunication", "GrpcCommunication"]