digitalkin 0.3.1.dev0__tar.gz → 0.3.1.dev2__tar.gz
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.
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/PKG-INFO +5 -5
- digitalkin-0.3.1.dev2/examples/modules/dynamic_setup_module.py +362 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/pyproject.toml +17 -18
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/__version__.py +1 -1
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/core/job_manager/taskiq_broker.py +1 -1
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/core/job_manager/taskiq_job_manager.py +1 -2
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/core/task_manager/base_task_manager.py +87 -12
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/core/task_manager/task_executor.py +103 -27
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/core/task_manager/task_session.py +75 -19
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/grpc_servers/module_servicer.py +17 -9
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/models/core/task_monitor.py +17 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/models/grpc_servers/models.py +4 -4
- digitalkin-0.3.1.dev2/src/digitalkin/models/module/module_types.py +393 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/modules/_base_module.py +66 -28
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/services/user_profile/grpc_user_profile.py +2 -2
- digitalkin-0.3.1.dev2/src/digitalkin/utils/__init__.py +29 -0
- digitalkin-0.3.1.dev2/src/digitalkin/utils/dynamic_schema.py +483 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin.egg-info/PKG-INFO +5 -5
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin.egg-info/SOURCES.txt +2 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin.egg-info/requires.txt +4 -4
- digitalkin-0.3.1.dev0/src/digitalkin/models/module/module_types.py +0 -109
- digitalkin-0.3.1.dev0/src/digitalkin/utils/__init__.py +0 -1
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/LICENSE +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/README.md +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/examples/base_server/__init__.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/examples/base_server/mock/__init__.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/examples/base_server/mock/mock_pb2.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/examples/base_server/mock/mock_pb2_grpc.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/examples/base_server/server_async_insecure.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/examples/base_server/server_async_secure.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/examples/base_server/server_sync_insecure.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/examples/base_server/server_sync_secure.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/examples/modules/__init__.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/examples/modules/cpu_intensive_module.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/examples/modules/minimal_llm_module.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/examples/modules/text_transform_module.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/examples/services/filesystem_module.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/examples/services/storage_module.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/setup.cfg +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/__init__.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/core/__init__.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/core/common/__init__.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/core/common/factories.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/core/job_manager/__init__.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/core/job_manager/base_job_manager.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/core/job_manager/single_job_manager.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/core/task_manager/__init__.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/core/task_manager/local_task_manager.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/core/task_manager/remote_task_manager.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/core/task_manager/surrealdb_repository.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/grpc_servers/__init__.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/grpc_servers/_base_server.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/grpc_servers/module_server.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/grpc_servers/registry_server.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/grpc_servers/registry_servicer.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/grpc_servers/utils/__init__.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/grpc_servers/utils/exceptions.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/grpc_servers/utils/grpc_client_wrapper.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/grpc_servers/utils/grpc_error_handler.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/logger.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/mixins/__init__.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/mixins/base_mixin.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/mixins/callback_mixin.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/mixins/chat_history_mixin.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/mixins/cost_mixin.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/mixins/file_history_mixin.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/mixins/filesystem_mixin.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/mixins/logger_mixin.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/mixins/storage_mixin.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/models/__init__.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/models/core/__init__.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/models/core/job_manager_models.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/models/grpc_servers/__init__.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/models/grpc_servers/types.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/models/module/__init__.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/models/module/module.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/models/module/module_context.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/models/services/__init__.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/models/services/cost.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/models/services/storage.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/modules/__init__.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/modules/archetype_module.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/modules/tool_module.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/modules/trigger_handler.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/py.typed +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/services/__init__.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/services/agent/__init__.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/services/agent/agent_strategy.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/services/agent/default_agent.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/services/base_strategy.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/services/cost/__init__.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/services/cost/cost_strategy.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/services/cost/default_cost.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/services/cost/grpc_cost.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/services/filesystem/__init__.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/services/filesystem/default_filesystem.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/services/filesystem/filesystem_strategy.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/services/filesystem/grpc_filesystem.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/services/identity/__init__.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/services/identity/default_identity.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/services/identity/identity_strategy.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/services/registry/__init__.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/services/registry/default_registry.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/services/registry/registry_strategy.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/services/services_config.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/services/services_models.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/services/setup/__init__.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/services/setup/default_setup.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/services/setup/grpc_setup.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/services/setup/setup_strategy.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/services/snapshot/__init__.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/services/snapshot/default_snapshot.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/services/snapshot/snapshot_strategy.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/services/storage/__init__.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/services/storage/default_storage.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/services/storage/grpc_storage.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/services/storage/storage_strategy.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/services/user_profile/__init__.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/services/user_profile/default_user_profile.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/services/user_profile/user_profile_strategy.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/utils/arg_parser.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/utils/development_mode_action.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/utils/llm_ready_schema.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/utils/package_discover.py +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin.egg-info/dependency_links.txt +0 -0
- {digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin.egg-info/top_level.txt +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: digitalkin
|
|
3
|
-
Version: 0.3.1.
|
|
3
|
+
Version: 0.3.1.dev2
|
|
4
4
|
Summary: SDK to build kin used in DigitalKin
|
|
5
5
|
Author-email: "DigitalKin.ai" <contact@digitalkin.ai>
|
|
6
6
|
License: Attribution-NonCommercial-ShareAlike 4.0 International
|
|
@@ -452,17 +452,17 @@ Classifier: License :: Other/Proprietary License
|
|
|
452
452
|
Requires-Python: >=3.10
|
|
453
453
|
Description-Content-Type: text/markdown
|
|
454
454
|
License-File: LICENSE
|
|
455
|
-
Requires-Dist: digitalkin-proto
|
|
455
|
+
Requires-Dist: digitalkin-proto==0.2.0.dev5
|
|
456
456
|
Requires-Dist: grpcio-health-checking>=1.76.0
|
|
457
457
|
Requires-Dist: grpcio-reflection>=1.76.0
|
|
458
458
|
Requires-Dist: grpcio-status>=1.76.0
|
|
459
|
-
Requires-Dist: pydantic>=2.12.
|
|
459
|
+
Requires-Dist: pydantic>=2.12.5
|
|
460
460
|
Requires-Dist: surrealdb>=1.0.6
|
|
461
461
|
Provides-Extra: taskiq
|
|
462
462
|
Requires-Dist: rstream>=0.40.0; extra == "taskiq"
|
|
463
|
-
Requires-Dist: taskiq-aio-pika>=0.
|
|
463
|
+
Requires-Dist: taskiq-aio-pika>=0.5.0; extra == "taskiq"
|
|
464
464
|
Requires-Dist: taskiq-redis>=1.1.2; extra == "taskiq"
|
|
465
|
-
Requires-Dist: taskiq[reload]>=0.
|
|
465
|
+
Requires-Dist: taskiq[reload]>=0.12.0; extra == "taskiq"
|
|
466
466
|
Dynamic: license-file
|
|
467
467
|
|
|
468
468
|
# DigitalKin Python SDK
|
|
@@ -0,0 +1,362 @@
|
|
|
1
|
+
"""Example module demonstrating dynamic schema fields in SetupModel.
|
|
2
|
+
|
|
3
|
+
This example shows how to use the Dynamic metadata class with async fetchers
|
|
4
|
+
to populate field schemas (like enums) at runtime. This is useful when the
|
|
5
|
+
available options come from external sources like databases or APIs.
|
|
6
|
+
|
|
7
|
+
Usage:
|
|
8
|
+
# Start the module server
|
|
9
|
+
python examples/modules/dynamic_setup_module.py
|
|
10
|
+
|
|
11
|
+
# Or import and use in your own code
|
|
12
|
+
from examples.modules.dynamic_setup_module import DynamicSetupModule
|
|
13
|
+
"""
|
|
14
|
+
|
|
15
|
+
import asyncio
|
|
16
|
+
import logging
|
|
17
|
+
from typing import Annotated, Any, ClassVar
|
|
18
|
+
|
|
19
|
+
from pydantic import BaseModel, Field
|
|
20
|
+
|
|
21
|
+
from digitalkin.models.module.module_context import ModuleContext
|
|
22
|
+
from digitalkin.models.module.module_types import DataModel, DataTrigger, SetupModel
|
|
23
|
+
from digitalkin.modules._base_module import BaseModule
|
|
24
|
+
from digitalkin.services.services_models import ServicesStrategy
|
|
25
|
+
from digitalkin.utils import Dynamic
|
|
26
|
+
|
|
27
|
+
# Configure logging
|
|
28
|
+
logging.basicConfig(
|
|
29
|
+
level=logging.DEBUG,
|
|
30
|
+
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
|
31
|
+
)
|
|
32
|
+
logger = logging.getLogger(__name__)
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
# =============================================================================
|
|
36
|
+
# Simulated External Services (replace with real implementations)
|
|
37
|
+
# =============================================================================
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
class MockModelRegistry:
|
|
41
|
+
"""Simulates an external model registry service.
|
|
42
|
+
|
|
43
|
+
In a real application, this would be a connection to a database,
|
|
44
|
+
API service, or configuration management system.
|
|
45
|
+
"""
|
|
46
|
+
|
|
47
|
+
_models: ClassVar[list[str]] = ["gpt-4", "gpt-4-turbo", "gpt-3.5-turbo"]
|
|
48
|
+
_languages: ClassVar[list[str]] = ["en", "fr", "de", "es", "it", "pt"]
|
|
49
|
+
|
|
50
|
+
@classmethod
|
|
51
|
+
async def fetch_available_models(cls) -> list[str]:
|
|
52
|
+
"""Fetch available models from the registry.
|
|
53
|
+
|
|
54
|
+
Simulates an async API call with a small delay.
|
|
55
|
+
"""
|
|
56
|
+
await asyncio.sleep(0.1) # Simulate network latency
|
|
57
|
+
logger.info("Fetched %d models from registry", len(cls._models))
|
|
58
|
+
return cls._models.copy()
|
|
59
|
+
|
|
60
|
+
@classmethod
|
|
61
|
+
async def fetch_supported_languages(cls) -> list[str]:
|
|
62
|
+
"""Fetch supported languages from the registry."""
|
|
63
|
+
await asyncio.sleep(0.05) # Simulate network latency
|
|
64
|
+
logger.info("Fetched %d languages from registry", len(cls._languages))
|
|
65
|
+
return cls._languages.copy()
|
|
66
|
+
|
|
67
|
+
@classmethod
|
|
68
|
+
def get_default_model(cls) -> str:
|
|
69
|
+
"""Get the default model (sync fetcher example)."""
|
|
70
|
+
return cls._models[0] if cls._models else "gpt-4"
|
|
71
|
+
|
|
72
|
+
|
|
73
|
+
# =============================================================================
|
|
74
|
+
# Dynamic Fetcher Functions
|
|
75
|
+
# =============================================================================
|
|
76
|
+
|
|
77
|
+
|
|
78
|
+
async def fetch_models() -> list[str]:
|
|
79
|
+
"""Async fetcher for available model names.
|
|
80
|
+
|
|
81
|
+
This function is called when SetupModel.get_clean_model(force=True)
|
|
82
|
+
is invoked, typically during module initialization or schema refresh.
|
|
83
|
+
"""
|
|
84
|
+
return await MockModelRegistry.fetch_available_models()
|
|
85
|
+
|
|
86
|
+
|
|
87
|
+
async def fetch_languages() -> list[str]:
|
|
88
|
+
"""Async fetcher for supported languages."""
|
|
89
|
+
return await MockModelRegistry.fetch_supported_languages()
|
|
90
|
+
|
|
91
|
+
|
|
92
|
+
def get_temperature_range() -> dict[str, float]:
|
|
93
|
+
"""Sync fetcher example returning min/max for temperature.
|
|
94
|
+
|
|
95
|
+
Demonstrates that fetchers can return any JSON-serializable value,
|
|
96
|
+
not just lists for enums.
|
|
97
|
+
"""
|
|
98
|
+
return {"minimum": 0.0, "maximum": 2.0}
|
|
99
|
+
|
|
100
|
+
|
|
101
|
+
# =============================================================================
|
|
102
|
+
# Setup Model with Dynamic Fields
|
|
103
|
+
# =============================================================================
|
|
104
|
+
|
|
105
|
+
|
|
106
|
+
class DynamicAgentSetup(SetupModel):
|
|
107
|
+
"""Setup model demonstrating dynamic schema fields.
|
|
108
|
+
|
|
109
|
+
Fields marked with Dynamic(...) will have their schema values
|
|
110
|
+
refreshed at runtime when get_clean_model(force=True) is called.
|
|
111
|
+
|
|
112
|
+
Attributes:
|
|
113
|
+
model_name: The LLM model to use. Enum values fetched from registry.
|
|
114
|
+
language: Output language. Enum values fetched dynamically.
|
|
115
|
+
temperature: Sampling temperature. Static field for comparison.
|
|
116
|
+
max_tokens: Maximum tokens to generate.
|
|
117
|
+
system_prompt: The system prompt for the model.
|
|
118
|
+
"""
|
|
119
|
+
|
|
120
|
+
# Dynamic field: enum values fetched asynchronously from model registry
|
|
121
|
+
model_name: Annotated[str, Dynamic(enum=fetch_models)] = Field(
|
|
122
|
+
default="gpt-4",
|
|
123
|
+
title="Model Name",
|
|
124
|
+
description="The LLM model to use for generation.",
|
|
125
|
+
json_schema_extra={
|
|
126
|
+
"config": True, # Shown in initial configuration
|
|
127
|
+
"ui:widget": "select",
|
|
128
|
+
},
|
|
129
|
+
)
|
|
130
|
+
|
|
131
|
+
# Dynamic field: language options fetched asynchronously
|
|
132
|
+
language: Annotated[str, Dynamic(enum=fetch_languages)] = Field(
|
|
133
|
+
default="en",
|
|
134
|
+
title="Output Language",
|
|
135
|
+
description="The language for generated responses.",
|
|
136
|
+
json_schema_extra={
|
|
137
|
+
"config": True,
|
|
138
|
+
"ui:widget": "select",
|
|
139
|
+
},
|
|
140
|
+
)
|
|
141
|
+
|
|
142
|
+
# Static field: no dynamic fetcher, values defined at class time
|
|
143
|
+
temperature: float = Field(
|
|
144
|
+
default=0.7,
|
|
145
|
+
ge=0.0,
|
|
146
|
+
le=2.0,
|
|
147
|
+
title="Temperature",
|
|
148
|
+
description="Controls randomness. Higher values = more creative.",
|
|
149
|
+
json_schema_extra={"config": True},
|
|
150
|
+
)
|
|
151
|
+
|
|
152
|
+
# Static field with hidden flag (runtime-only, not in initial config)
|
|
153
|
+
max_tokens: int = Field(
|
|
154
|
+
default=1024,
|
|
155
|
+
ge=1,
|
|
156
|
+
le=4096,
|
|
157
|
+
title="Max Tokens",
|
|
158
|
+
description="Maximum tokens in the response.",
|
|
159
|
+
json_schema_extra={"hidden": True},
|
|
160
|
+
)
|
|
161
|
+
|
|
162
|
+
# Static field without any special flags
|
|
163
|
+
system_prompt: str = Field(
|
|
164
|
+
default="You are a helpful assistant.",
|
|
165
|
+
title="System Prompt",
|
|
166
|
+
description="The system prompt defining assistant behavior.",
|
|
167
|
+
)
|
|
168
|
+
|
|
169
|
+
|
|
170
|
+
# =============================================================================
|
|
171
|
+
# Input/Output Models (Using DataModel/DataTrigger pattern)
|
|
172
|
+
# =============================================================================
|
|
173
|
+
|
|
174
|
+
|
|
175
|
+
class MessageInputTrigger(DataTrigger):
|
|
176
|
+
"""Message input trigger following DigitalKin DataTrigger pattern.
|
|
177
|
+
|
|
178
|
+
The protocol field determines which trigger handler processes this input.
|
|
179
|
+
"""
|
|
180
|
+
|
|
181
|
+
protocol: str = "message"
|
|
182
|
+
content: str = Field(default="", description="The user message content.")
|
|
183
|
+
|
|
184
|
+
|
|
185
|
+
class DynamicModuleInput(DataModel[MessageInputTrigger]):
|
|
186
|
+
"""Input model following DigitalKin DataModel pattern.
|
|
187
|
+
|
|
188
|
+
Wraps the trigger in a root field with optional annotations.
|
|
189
|
+
"""
|
|
190
|
+
|
|
191
|
+
root: MessageInputTrigger = Field(default_factory=MessageInputTrigger)
|
|
192
|
+
|
|
193
|
+
|
|
194
|
+
class MessageOutputTrigger(DataTrigger):
|
|
195
|
+
"""Message output trigger following DigitalKin DataTrigger pattern."""
|
|
196
|
+
|
|
197
|
+
protocol: str = "message"
|
|
198
|
+
content: str = Field(default="", description="The generated response.")
|
|
199
|
+
model_used: str = Field(default="", description="The model that generated this response.")
|
|
200
|
+
language: str = Field(default="", description="The output language.")
|
|
201
|
+
|
|
202
|
+
|
|
203
|
+
class DynamicModuleOutput(DataModel[MessageOutputTrigger]):
|
|
204
|
+
"""Output model following DigitalKin DataModel pattern."""
|
|
205
|
+
|
|
206
|
+
root: MessageOutputTrigger = Field(default_factory=MessageOutputTrigger)
|
|
207
|
+
|
|
208
|
+
|
|
209
|
+
class DynamicModuleSecret(BaseModel):
|
|
210
|
+
"""Secret model (empty for this example)."""
|
|
211
|
+
|
|
212
|
+
pass
|
|
213
|
+
|
|
214
|
+
|
|
215
|
+
# =============================================================================
|
|
216
|
+
# Module Implementation
|
|
217
|
+
# =============================================================================
|
|
218
|
+
|
|
219
|
+
|
|
220
|
+
class DynamicSetupModule(
|
|
221
|
+
BaseModule[
|
|
222
|
+
DynamicModuleInput,
|
|
223
|
+
DynamicModuleOutput,
|
|
224
|
+
DynamicAgentSetup,
|
|
225
|
+
DynamicModuleSecret,
|
|
226
|
+
]
|
|
227
|
+
):
|
|
228
|
+
"""Example module demonstrating dynamic schema in SetupModel.
|
|
229
|
+
|
|
230
|
+
This module shows how to:
|
|
231
|
+
1. Define setup fields with Dynamic() metadata for runtime enum fetching
|
|
232
|
+
2. Mix static and dynamic fields in the same SetupModel
|
|
233
|
+
3. Use async fetchers that simulate external service calls
|
|
234
|
+
4. Follow DigitalKin's DataModel/DataTrigger pattern for I/O
|
|
235
|
+
|
|
236
|
+
The key integration point is in the gRPC servicer, which calls
|
|
237
|
+
SetupModel.get_clean_model(force=True) to refresh dynamic values
|
|
238
|
+
before returning schema information to clients.
|
|
239
|
+
"""
|
|
240
|
+
|
|
241
|
+
name = "DynamicSetupModule"
|
|
242
|
+
description = "Demonstrates dynamic schema fields in module setup"
|
|
243
|
+
|
|
244
|
+
# Schema format definitions
|
|
245
|
+
input_format = DynamicModuleInput
|
|
246
|
+
output_format = DynamicModuleOutput
|
|
247
|
+
setup_format = DynamicAgentSetup
|
|
248
|
+
secret_format = DynamicModuleSecret
|
|
249
|
+
|
|
250
|
+
# Module metadata
|
|
251
|
+
metadata: ClassVar[dict[str, Any]] = {
|
|
252
|
+
"name": "DynamicSetupModule",
|
|
253
|
+
"description": "Example module with dynamic setup schema",
|
|
254
|
+
"version": "1.0.0",
|
|
255
|
+
"tags": ["example", "dynamic-schema"],
|
|
256
|
+
}
|
|
257
|
+
|
|
258
|
+
# Services configuration (empty for this example)
|
|
259
|
+
services_config_strategies: ClassVar[dict[str, ServicesStrategy | None]] = {}
|
|
260
|
+
services_config_params: ClassVar[dict[str, dict[str, Any | None] | None]] = {}
|
|
261
|
+
|
|
262
|
+
async def initialize(self, context: ModuleContext, setup_data: DynamicAgentSetup) -> None:
|
|
263
|
+
"""Initialize the module with setup data.
|
|
264
|
+
|
|
265
|
+
Args:
|
|
266
|
+
context: The module context with services and session info.
|
|
267
|
+
setup_data: The validated setup configuration.
|
|
268
|
+
"""
|
|
269
|
+
logger.info(
|
|
270
|
+
"Initializing DynamicSetupModule with model=%s, language=%s",
|
|
271
|
+
setup_data.model_name,
|
|
272
|
+
setup_data.language,
|
|
273
|
+
)
|
|
274
|
+
self.setup = setup_data
|
|
275
|
+
|
|
276
|
+
async def cleanup(self) -> None:
|
|
277
|
+
"""Clean up resources."""
|
|
278
|
+
logger.info("Cleaning up DynamicSetupModule")
|
|
279
|
+
|
|
280
|
+
|
|
281
|
+
# =============================================================================
|
|
282
|
+
# Demonstration Script
|
|
283
|
+
# =============================================================================
|
|
284
|
+
|
|
285
|
+
|
|
286
|
+
async def demonstrate_dynamic_schema() -> None:
|
|
287
|
+
"""Demonstrate the dynamic schema functionality."""
|
|
288
|
+
print("=" * 60)
|
|
289
|
+
print("Dynamic Schema Demonstration")
|
|
290
|
+
print("=" * 60)
|
|
291
|
+
|
|
292
|
+
# 1. Show schema WITHOUT force (dynamic fields not resolved)
|
|
293
|
+
print("\n1. Schema without force=True (fetchers NOT called):")
|
|
294
|
+
print("-" * 40)
|
|
295
|
+
|
|
296
|
+
model_no_force = await DynamicAgentSetup.get_clean_model(
|
|
297
|
+
config_fields=True,
|
|
298
|
+
hidden_fields=False,
|
|
299
|
+
force=False,
|
|
300
|
+
)
|
|
301
|
+
schema_no_force = model_no_force.model_json_schema()
|
|
302
|
+
|
|
303
|
+
# Check if enum is present
|
|
304
|
+
model_name_schema = schema_no_force.get("properties", {}).get("model_name", {})
|
|
305
|
+
print(f"model_name has enum: {'enum' in model_name_schema}")
|
|
306
|
+
if "enum" in model_name_schema:
|
|
307
|
+
print(f" enum values: {model_name_schema['enum']}")
|
|
308
|
+
|
|
309
|
+
# 2. Show schema WITH force (dynamic fields resolved)
|
|
310
|
+
print("\n2. Schema with force=True (fetchers called):")
|
|
311
|
+
print("-" * 40)
|
|
312
|
+
|
|
313
|
+
model_with_force = await DynamicAgentSetup.get_clean_model(
|
|
314
|
+
config_fields=True,
|
|
315
|
+
hidden_fields=False,
|
|
316
|
+
force=True,
|
|
317
|
+
)
|
|
318
|
+
schema_with_force = model_with_force.model_json_schema()
|
|
319
|
+
|
|
320
|
+
# Check enum values after force
|
|
321
|
+
model_name_schema = schema_with_force.get("properties", {}).get("model_name", {})
|
|
322
|
+
print(f"model_name has enum: {'enum' in model_name_schema}")
|
|
323
|
+
if "enum" in model_name_schema:
|
|
324
|
+
print(f" enum values: {model_name_schema['enum']}")
|
|
325
|
+
|
|
326
|
+
language_schema = schema_with_force.get("properties", {}).get("language", {})
|
|
327
|
+
print(f"language has enum: {'enum' in language_schema}")
|
|
328
|
+
if "enum" in language_schema:
|
|
329
|
+
print(f" enum values: {language_schema['enum']}")
|
|
330
|
+
|
|
331
|
+
# 3. Show that static json_schema_extra is preserved
|
|
332
|
+
print("\n3. Static json_schema_extra preserved:")
|
|
333
|
+
print("-" * 40)
|
|
334
|
+
print(f"model_name ui:widget: {model_name_schema.get('ui:widget', 'NOT FOUND')}")
|
|
335
|
+
|
|
336
|
+
# 4. Show field filtering
|
|
337
|
+
print("\n4. Field filtering demonstration:")
|
|
338
|
+
print("-" * 40)
|
|
339
|
+
|
|
340
|
+
# Config fields only (hidden excluded)
|
|
341
|
+
config_model = await DynamicAgentSetup.get_clean_model(
|
|
342
|
+
config_fields=True,
|
|
343
|
+
hidden_fields=False,
|
|
344
|
+
force=False,
|
|
345
|
+
)
|
|
346
|
+
print(f"Config fields (hidden=False): {list(config_model.model_fields.keys())}")
|
|
347
|
+
|
|
348
|
+
# All fields including hidden
|
|
349
|
+
all_model = await DynamicAgentSetup.get_clean_model(
|
|
350
|
+
config_fields=True,
|
|
351
|
+
hidden_fields=True,
|
|
352
|
+
force=False,
|
|
353
|
+
)
|
|
354
|
+
print(f"All fields (hidden=True): {list(all_model.model_fields.keys())}")
|
|
355
|
+
|
|
356
|
+
print("\n" + "=" * 60)
|
|
357
|
+
print("Demonstration complete!")
|
|
358
|
+
print("=" * 60)
|
|
359
|
+
|
|
360
|
+
|
|
361
|
+
if __name__ == "__main__":
|
|
362
|
+
asyncio.run(demonstrate_dynamic_schema())
|
|
@@ -12,7 +12,7 @@
|
|
|
12
12
|
|
|
13
13
|
keywords = [ "digitalkin", "kin", "agent", "gprc", "sdk" ]
|
|
14
14
|
|
|
15
|
-
version = "0.3.1.
|
|
15
|
+
version = "0.3.1.dev2"
|
|
16
16
|
classifiers = [
|
|
17
17
|
"Development Status :: 3 - Alpha",
|
|
18
18
|
"Intended Audience :: Developers",
|
|
@@ -28,20 +28,20 @@
|
|
|
28
28
|
]
|
|
29
29
|
|
|
30
30
|
dependencies = [
|
|
31
|
-
"digitalkin-proto
|
|
31
|
+
"digitalkin-proto==0.2.0.dev5",
|
|
32
32
|
"grpcio-health-checking>=1.76.0",
|
|
33
33
|
"grpcio-reflection>=1.76.0",
|
|
34
34
|
"grpcio-status>=1.76.0",
|
|
35
|
-
"pydantic>=2.12.
|
|
35
|
+
"pydantic>=2.12.5",
|
|
36
36
|
"surrealdb>=1.0.6",
|
|
37
37
|
]
|
|
38
38
|
|
|
39
39
|
[project.optional-dependencies]
|
|
40
40
|
taskiq = [
|
|
41
41
|
"rstream>=0.40.0",
|
|
42
|
-
"taskiq-aio-pika>=0.
|
|
42
|
+
"taskiq-aio-pika>=0.5.0",
|
|
43
43
|
"taskiq-redis>=1.1.2",
|
|
44
|
-
"taskiq[reload]>=0.
|
|
44
|
+
"taskiq[reload]>=0.12.0",
|
|
45
45
|
]
|
|
46
46
|
|
|
47
47
|
[project.urls]
|
|
@@ -55,18 +55,18 @@
|
|
|
55
55
|
|
|
56
56
|
[dependency-groups]
|
|
57
57
|
dev = [
|
|
58
|
-
"typos>=1.
|
|
59
|
-
"ruff>=0.14.
|
|
60
|
-
"mypy>=1.
|
|
58
|
+
"typos>=1.40.0",
|
|
59
|
+
"ruff>=0.14.7",
|
|
60
|
+
"mypy>=1.19.0",
|
|
61
61
|
"pyright>=1.1.407",
|
|
62
|
-
"pre-commit>=4.
|
|
62
|
+
"pre-commit>=4.5.0",
|
|
63
63
|
"bump-my-version>=1.2.4",
|
|
64
64
|
"build>=1.3.0",
|
|
65
65
|
"twine>=6.2.0",
|
|
66
66
|
"cryptography>=46.0.3",
|
|
67
67
|
]
|
|
68
68
|
tests = [
|
|
69
|
-
"freezegun>=1.5.
|
|
69
|
+
"freezegun>=1.5.5",
|
|
70
70
|
"hdrhistogram>=0.10.3",
|
|
71
71
|
"grpcio-testing>=1.76.0",
|
|
72
72
|
"psutil>=7.1.3",
|
|
@@ -75,17 +75,17 @@
|
|
|
75
75
|
"pytest-cov>=7.0.0",
|
|
76
76
|
"pytest-html==4.1.1",
|
|
77
77
|
"pytest-json-report==1.5.0",
|
|
78
|
-
"pytest-timeout>=2.
|
|
78
|
+
"pytest-timeout>=2.4.0",
|
|
79
79
|
]
|
|
80
80
|
docs = [
|
|
81
81
|
"mike>=2.1.3",
|
|
82
|
-
"markdown-callouts>=0.4",
|
|
82
|
+
"markdown-callouts>=0.4.0",
|
|
83
83
|
"markdown-exec>=1.12.1",
|
|
84
|
-
"mkdocs>=1.6",
|
|
84
|
+
"mkdocs>=1.6.1",
|
|
85
85
|
"mkdocs-coverage>=2.0.0",
|
|
86
|
-
"mkdocs-llmstxt>=0.
|
|
87
|
-
"mkdocs-redirects>=1.2",
|
|
88
|
-
"mkdocstrings>=0.
|
|
86
|
+
"mkdocs-llmstxt>=0.5.0",
|
|
87
|
+
"mkdocs-redirects>=1.2.2",
|
|
88
|
+
"mkdocstrings>=1.0.0",
|
|
89
89
|
"griffe-inherited-docstrings>=1.1.2",
|
|
90
90
|
"mkdocs-autorefs>=1.4.3",
|
|
91
91
|
"mkdocs-awesome-pages-plugin>=2.10.1",
|
|
@@ -97,8 +97,7 @@
|
|
|
97
97
|
"mkdocs-material[imaging]>=9.7.0",
|
|
98
98
|
"mkdocs-minify-plugin>=0.8.0",
|
|
99
99
|
"mkdocs-section-index>=0.3.10",
|
|
100
|
-
"mkdocstrings>=0.
|
|
101
|
-
"mkdocstrings-python>=1.19.0",
|
|
100
|
+
"mkdocstrings-python>=2.0.0",
|
|
102
101
|
"mkdocs-open-in-new-tab>=1.0.8",
|
|
103
102
|
"tomli>=2.3.0",
|
|
104
103
|
]
|
{digitalkin-0.3.1.dev0 → digitalkin-0.3.1.dev2}/src/digitalkin/core/job_manager/taskiq_broker.py
RENAMED
|
@@ -208,7 +208,7 @@ async def run_start_module(
|
|
|
208
208
|
# Reconstruct Pydantic models from dicts for type safety
|
|
209
209
|
try:
|
|
210
210
|
input_model = module_class.create_input_model(input_data)
|
|
211
|
-
setup_model = module_class.create_setup_model(setup_data)
|
|
211
|
+
setup_model = await module_class.create_setup_model(setup_data)
|
|
212
212
|
except Exception as e:
|
|
213
213
|
logger.error("Failed to reconstruct models for job %s: %s", job_id, e, exc_info=True)
|
|
214
214
|
raise
|
|
@@ -140,7 +140,7 @@ class TaskiqJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
|
|
|
140
140
|
services_mode: ServicesMode,
|
|
141
141
|
default_timeout: float = 10.0,
|
|
142
142
|
max_concurrent_tasks: int = 100,
|
|
143
|
-
stream_timeout: float =
|
|
143
|
+
stream_timeout: float = 30.0,
|
|
144
144
|
) -> None:
|
|
145
145
|
"""Initialize the Taskiq job manager.
|
|
146
146
|
|
|
@@ -298,7 +298,6 @@ class TaskiqJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
|
|
|
298
298
|
while True:
|
|
299
299
|
try:
|
|
300
300
|
# Block for first item with timeout to allow termination checks
|
|
301
|
-
# Configurable timeout (default 15s) to account for distributed system latencies
|
|
302
301
|
item = await asyncio.wait_for(queue.get(), timeout=self.stream_timeout)
|
|
303
302
|
queue.task_done()
|
|
304
303
|
yield item
|
|
@@ -11,6 +11,7 @@ from typing import Any
|
|
|
11
11
|
from digitalkin.core.task_manager.surrealdb_repository import SurrealDBConnection
|
|
12
12
|
from digitalkin.core.task_manager.task_session import TaskSession
|
|
13
13
|
from digitalkin.logger import logger
|
|
14
|
+
from digitalkin.models.core.task_monitor import CancellationReason
|
|
14
15
|
from digitalkin.modules._base_module import BaseModule
|
|
15
16
|
|
|
16
17
|
|
|
@@ -84,13 +85,32 @@ class BaseTaskManager(ABC):
|
|
|
84
85
|
task_id: The ID of the task to clean up
|
|
85
86
|
mission_id: The ID of the mission associated with the task
|
|
86
87
|
"""
|
|
88
|
+
session = self.tasks_sessions.get(task_id)
|
|
89
|
+
cancellation_reason = session.cancellation_reason.value if session else "no_session"
|
|
90
|
+
final_status = session.status.value if session else "unknown"
|
|
91
|
+
|
|
87
92
|
logger.debug(
|
|
88
|
-
"Cleaning up resources
|
|
93
|
+
"Cleaning up resources",
|
|
94
|
+
extra={
|
|
95
|
+
"mission_id": mission_id,
|
|
96
|
+
"task_id": task_id,
|
|
97
|
+
"final_status": final_status,
|
|
98
|
+
"cancellation_reason": cancellation_reason,
|
|
99
|
+
},
|
|
89
100
|
)
|
|
90
|
-
|
|
91
|
-
|
|
101
|
+
|
|
102
|
+
if session:
|
|
92
103
|
await session.cleanup()
|
|
93
104
|
self.tasks_sessions.pop(task_id, None)
|
|
105
|
+
logger.debug(
|
|
106
|
+
"Task session cleanup completed",
|
|
107
|
+
extra={
|
|
108
|
+
"mission_id": mission_id,
|
|
109
|
+
"task_id": task_id,
|
|
110
|
+
"final_status": final_status,
|
|
111
|
+
"cancellation_reason": cancellation_reason,
|
|
112
|
+
},
|
|
113
|
+
)
|
|
94
114
|
|
|
95
115
|
self.tasks.pop(task_id, None)
|
|
96
116
|
|
|
@@ -261,10 +281,21 @@ class BaseTaskManager(ABC):
|
|
|
261
281
|
)
|
|
262
282
|
|
|
263
283
|
except asyncio.TimeoutError:
|
|
284
|
+
# Set timeout as cancellation reason
|
|
285
|
+
if task_id in self.tasks_sessions:
|
|
286
|
+
session = self.tasks_sessions[task_id]
|
|
287
|
+
if session.cancellation_reason == CancellationReason.UNKNOWN:
|
|
288
|
+
session.cancellation_reason = CancellationReason.TIMEOUT
|
|
289
|
+
|
|
264
290
|
logger.warning(
|
|
265
291
|
"Graceful cancellation timed out for task: '%s', forcing cancellation",
|
|
266
292
|
task_id,
|
|
267
|
-
extra={
|
|
293
|
+
extra={
|
|
294
|
+
"mission_id": mission_id,
|
|
295
|
+
"task_id": task_id,
|
|
296
|
+
"timeout": timeout,
|
|
297
|
+
"cancellation_reason": CancellationReason.TIMEOUT.value,
|
|
298
|
+
},
|
|
268
299
|
)
|
|
269
300
|
|
|
270
301
|
# Phase 2: Force cancellation
|
|
@@ -272,8 +303,16 @@ class BaseTaskManager(ABC):
|
|
|
272
303
|
with contextlib.suppress(asyncio.CancelledError):
|
|
273
304
|
await task
|
|
274
305
|
|
|
275
|
-
logger.warning(
|
|
276
|
-
|
|
306
|
+
logger.warning(
|
|
307
|
+
"Task force-cancelled: '%s', reason: %s",
|
|
308
|
+
task_id,
|
|
309
|
+
CancellationReason.TIMEOUT.value,
|
|
310
|
+
extra={
|
|
311
|
+
"mission_id": mission_id,
|
|
312
|
+
"task_id": task_id,
|
|
313
|
+
"cancellation_reason": CancellationReason.TIMEOUT.value,
|
|
314
|
+
},
|
|
315
|
+
)
|
|
277
316
|
return True
|
|
278
317
|
|
|
279
318
|
except Exception as e:
|
|
@@ -283,10 +322,9 @@ class BaseTaskManager(ABC):
|
|
|
283
322
|
extra={"mission_id": mission_id, "task_id": task_id, "error": str(e)},
|
|
284
323
|
exc_info=True,
|
|
285
324
|
)
|
|
286
|
-
await self._cleanup_task(task_id, mission_id)
|
|
287
325
|
return False
|
|
288
|
-
|
|
289
|
-
|
|
326
|
+
finally:
|
|
327
|
+
await self._cleanup_task(task_id, mission_id)
|
|
290
328
|
return True
|
|
291
329
|
|
|
292
330
|
async def clean_session(self, task_id: str, mission_id: str) -> bool:
|
|
@@ -382,7 +420,16 @@ class BaseTaskManager(ABC):
|
|
|
382
420
|
results: dict[str, bool | BaseException] = {}
|
|
383
421
|
for task_id, result in zip(task_ids, results_list):
|
|
384
422
|
if isinstance(result, Exception):
|
|
385
|
-
logger.error(
|
|
423
|
+
logger.error(
|
|
424
|
+
"Exception cancelling task: '%s', error: %s",
|
|
425
|
+
task_id,
|
|
426
|
+
result,
|
|
427
|
+
extra={
|
|
428
|
+
"mission_id": mission_id,
|
|
429
|
+
"task_id": task_id,
|
|
430
|
+
"error": str(result),
|
|
431
|
+
},
|
|
432
|
+
)
|
|
386
433
|
results[task_id] = False
|
|
387
434
|
else:
|
|
388
435
|
results[task_id] = result
|
|
@@ -403,6 +450,21 @@ class BaseTaskManager(ABC):
|
|
|
403
450
|
)
|
|
404
451
|
|
|
405
452
|
self._shutdown_event.set()
|
|
453
|
+
|
|
454
|
+
# Mark all sessions with shutdown reason before cancellation
|
|
455
|
+
for task_id, session in self.tasks_sessions.items():
|
|
456
|
+
if session.cancellation_reason == CancellationReason.UNKNOWN:
|
|
457
|
+
session.cancellation_reason = CancellationReason.SHUTDOWN
|
|
458
|
+
logger.debug(
|
|
459
|
+
"Marking task for shutdown: '%s'",
|
|
460
|
+
task_id,
|
|
461
|
+
extra={
|
|
462
|
+
"mission_id": mission_id,
|
|
463
|
+
"task_id": task_id,
|
|
464
|
+
"cancellation_reason": CancellationReason.SHUTDOWN.value,
|
|
465
|
+
},
|
|
466
|
+
)
|
|
467
|
+
|
|
406
468
|
results = await self.cancel_all_tasks(mission_id, timeout)
|
|
407
469
|
|
|
408
470
|
failed_tasks = [task_id for task_id, success in results.items() if not success]
|
|
@@ -411,13 +473,26 @@ class BaseTaskManager(ABC):
|
|
|
411
473
|
"Failed to cancel %d tasks during shutdown: %s",
|
|
412
474
|
len(failed_tasks),
|
|
413
475
|
failed_tasks,
|
|
414
|
-
extra={
|
|
476
|
+
extra={
|
|
477
|
+
"mission_id": mission_id,
|
|
478
|
+
"failed_tasks": failed_tasks,
|
|
479
|
+
"failed_count": len(failed_tasks),
|
|
480
|
+
"cancellation_reason": CancellationReason.SHUTDOWN.value,
|
|
481
|
+
},
|
|
415
482
|
)
|
|
416
483
|
|
|
417
484
|
# Clean up any remaining sessions (in case cancellation didn't clean them)
|
|
418
485
|
remaining_sessions = list(self.tasks_sessions.keys())
|
|
419
486
|
if remaining_sessions:
|
|
420
|
-
logger.info(
|
|
487
|
+
logger.info(
|
|
488
|
+
"Cleaning up %d remaining task sessions after shutdown",
|
|
489
|
+
len(remaining_sessions),
|
|
490
|
+
extra={
|
|
491
|
+
"mission_id": mission_id,
|
|
492
|
+
"remaining_sessions": remaining_sessions,
|
|
493
|
+
"remaining_count": len(remaining_sessions),
|
|
494
|
+
},
|
|
495
|
+
)
|
|
421
496
|
cleanup_coros = [self._cleanup_task(task_id, mission_id) for task_id in remaining_sessions]
|
|
422
497
|
await asyncio.gather(*cleanup_coros, return_exceptions=True)
|
|
423
498
|
|