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.
- base_server/server_async_insecure.py +6 -5
- base_server/server_async_secure.py +6 -5
- base_server/server_sync_insecure.py +5 -4
- base_server/server_sync_secure.py +5 -4
- digitalkin/__version__.py +1 -1
- digitalkin/core/job_manager/base_job_manager.py +1 -1
- digitalkin/core/job_manager/single_job_manager.py +78 -36
- digitalkin/core/job_manager/taskiq_broker.py +7 -6
- digitalkin/core/job_manager/taskiq_job_manager.py +9 -5
- digitalkin/core/task_manager/base_task_manager.py +3 -1
- digitalkin/core/task_manager/surrealdb_repository.py +29 -7
- digitalkin/core/task_manager/task_executor.py +46 -12
- digitalkin/core/task_manager/task_session.py +132 -102
- digitalkin/grpc_servers/module_server.py +95 -171
- digitalkin/grpc_servers/module_servicer.py +121 -19
- digitalkin/grpc_servers/utils/grpc_client_wrapper.py +36 -10
- digitalkin/grpc_servers/utils/utility_schema_extender.py +106 -0
- digitalkin/models/__init__.py +1 -1
- digitalkin/models/core/job_manager_models.py +0 -8
- digitalkin/models/core/task_monitor.py +23 -1
- digitalkin/models/grpc_servers/models.py +95 -8
- digitalkin/models/module/__init__.py +26 -13
- digitalkin/models/module/base_types.py +61 -0
- digitalkin/models/module/module_context.py +279 -13
- digitalkin/models/module/module_types.py +28 -392
- digitalkin/models/module/setup_types.py +547 -0
- digitalkin/models/module/tool_cache.py +230 -0
- digitalkin/models/module/tool_reference.py +160 -0
- digitalkin/models/module/utility.py +167 -0
- digitalkin/models/services/cost.py +22 -1
- digitalkin/models/services/registry.py +77 -0
- digitalkin/modules/__init__.py +5 -1
- digitalkin/modules/_base_module.py +188 -63
- digitalkin/modules/archetype_module.py +6 -1
- digitalkin/modules/tool_module.py +6 -1
- digitalkin/modules/triggers/__init__.py +8 -0
- digitalkin/modules/triggers/healthcheck_ping_trigger.py +45 -0
- digitalkin/modules/triggers/healthcheck_services_trigger.py +63 -0
- digitalkin/modules/triggers/healthcheck_status_trigger.py +52 -0
- digitalkin/services/__init__.py +4 -0
- digitalkin/services/communication/__init__.py +7 -0
- digitalkin/services/communication/communication_strategy.py +87 -0
- digitalkin/services/communication/default_communication.py +104 -0
- digitalkin/services/communication/grpc_communication.py +264 -0
- digitalkin/services/cost/cost_strategy.py +36 -14
- digitalkin/services/cost/default_cost.py +61 -1
- digitalkin/services/cost/grpc_cost.py +98 -2
- digitalkin/services/filesystem/grpc_filesystem.py +9 -2
- digitalkin/services/registry/__init__.py +22 -1
- digitalkin/services/registry/default_registry.py +156 -4
- digitalkin/services/registry/exceptions.py +47 -0
- digitalkin/services/registry/grpc_registry.py +382 -0
- digitalkin/services/registry/registry_models.py +15 -0
- digitalkin/services/registry/registry_strategy.py +106 -4
- digitalkin/services/services_config.py +25 -3
- digitalkin/services/services_models.py +5 -1
- digitalkin/services/setup/default_setup.py +1 -1
- digitalkin/services/setup/grpc_setup.py +1 -1
- digitalkin/services/storage/grpc_storage.py +1 -1
- digitalkin/services/user_profile/__init__.py +11 -0
- digitalkin/services/user_profile/grpc_user_profile.py +2 -2
- digitalkin/services/user_profile/user_profile_strategy.py +0 -15
- digitalkin/utils/__init__.py +15 -3
- digitalkin/utils/conditional_schema.py +260 -0
- digitalkin/utils/dynamic_schema.py +4 -0
- digitalkin/utils/schema_splitter.py +290 -0
- {digitalkin-0.3.1.dev2.dist-info → digitalkin-0.3.2a3.dist-info}/METADATA +12 -12
- digitalkin-0.3.2a3.dist-info/RECORD +144 -0
- {digitalkin-0.3.1.dev2.dist-info → digitalkin-0.3.2a3.dist-info}/WHEEL +1 -1
- {digitalkin-0.3.1.dev2.dist-info → digitalkin-0.3.2a3.dist-info}/top_level.txt +1 -0
- modules/archetype_with_tools_module.py +232 -0
- modules/cpu_intensive_module.py +1 -1
- modules/dynamic_setup_module.py +5 -29
- modules/minimal_llm_module.py +1 -1
- modules/text_transform_module.py +1 -1
- monitoring/digitalkin_observability/__init__.py +46 -0
- monitoring/digitalkin_observability/http_server.py +150 -0
- monitoring/digitalkin_observability/interceptors.py +176 -0
- monitoring/digitalkin_observability/metrics.py +201 -0
- monitoring/digitalkin_observability/prometheus.py +137 -0
- monitoring/tests/test_metrics.py +172 -0
- services/filesystem_module.py +7 -5
- services/storage_module.py +4 -2
- digitalkin/grpc_servers/registry_server.py +0 -65
- digitalkin/grpc_servers/registry_servicer.py +0 -456
- digitalkin-0.3.1.dev2.dist-info/RECORD +0 -119
- {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.
|
|
18
|
-
from digitalkin.models.
|
|
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
|
-
|
|
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
|
|
145
|
-
|
|
146
|
-
|
|
147
|
-
|
|
148
|
-
|
|
149
|
-
|
|
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
|
|
166
|
-
|
|
167
|
-
|
|
168
|
-
|
|
169
|
-
|
|
170
|
-
|
|
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
|
-
|
|
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
|
-
|
|
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
|
|
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
|
|
364
|
-
setup_data
|
|
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
|
-
|
|
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(
|
|
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(
|
|
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
|
-
"""
|
|
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
|
-
|
|
475
|
-
|
|
476
|
-
|
|
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
|
|
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
|
|
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
|
|
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)
|
digitalkin/services/__init__.py
CHANGED
|
@@ -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"]
|