digitalkin 0.2.13__tar.gz → 0.2.15__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.
Files changed (95) hide show
  1. {digitalkin-0.2.13 → digitalkin-0.2.15}/PKG-INFO +14 -14
  2. {digitalkin-0.2.13 → digitalkin-0.2.15}/examples/modules/cpu_intensive_module.py +11 -2
  3. {digitalkin-0.2.13 → digitalkin-0.2.15}/examples/modules/minimal_llm_module.py +41 -3
  4. {digitalkin-0.2.13 → digitalkin-0.2.15}/examples/modules/text_transform_module.py +0 -1
  5. digitalkin-0.2.15/examples/services/filesystem_module.py +198 -0
  6. {digitalkin-0.2.13/examples/modules → digitalkin-0.2.15/examples/services}/storage_module.py +20 -7
  7. {digitalkin-0.2.13 → digitalkin-0.2.15}/pyproject.toml +41 -17
  8. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/__version__.py +1 -1
  9. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/grpc_servers/module_servicer.py +93 -1
  10. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/grpc_servers/utils/exceptions.py +4 -0
  11. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/models/module/__init__.py +2 -1
  12. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/models/module/module_types.py +1 -0
  13. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/modules/_base_module.py +80 -2
  14. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/modules/archetype_module.py +11 -1
  15. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/modules/job_manager/base_job_manager.py +45 -3
  16. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/modules/job_manager/single_job_manager.py +70 -3
  17. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/modules/job_manager/taskiq_broker.py +42 -1
  18. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/modules/job_manager/taskiq_job_manager.py +90 -17
  19. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/modules/tool_module.py +2 -1
  20. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/modules/trigger_module.py +3 -1
  21. digitalkin-0.2.15/src/digitalkin/services/filesystem/default_filesystem.py +405 -0
  22. digitalkin-0.2.15/src/digitalkin/services/filesystem/filesystem_strategy.py +206 -0
  23. digitalkin-0.2.15/src/digitalkin/services/filesystem/grpc_filesystem.py +322 -0
  24. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/services/setup/grpc_setup.py +1 -0
  25. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin.egg-info/PKG-INFO +14 -14
  26. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin.egg-info/SOURCES.txt +2 -1
  27. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin.egg-info/requires.txt +13 -13
  28. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin.egg-info/top_level.txt +1 -0
  29. digitalkin-0.2.13/src/digitalkin/services/filesystem/default_filesystem.py +0 -208
  30. digitalkin-0.2.13/src/digitalkin/services/filesystem/filesystem_strategy.py +0 -71
  31. digitalkin-0.2.13/src/digitalkin/services/filesystem/grpc_filesystem.py +0 -214
  32. {digitalkin-0.2.13 → digitalkin-0.2.15}/LICENSE +0 -0
  33. {digitalkin-0.2.13 → digitalkin-0.2.15}/README.md +0 -0
  34. {digitalkin-0.2.13 → digitalkin-0.2.15}/examples/base_server/__init__.py +0 -0
  35. {digitalkin-0.2.13 → digitalkin-0.2.15}/examples/base_server/mock/__init__.py +0 -0
  36. {digitalkin-0.2.13 → digitalkin-0.2.15}/examples/base_server/mock/mock_pb2.py +0 -0
  37. {digitalkin-0.2.13 → digitalkin-0.2.15}/examples/base_server/mock/mock_pb2_grpc.py +0 -0
  38. {digitalkin-0.2.13 → digitalkin-0.2.15}/examples/base_server/server_async_insecure.py +0 -0
  39. {digitalkin-0.2.13 → digitalkin-0.2.15}/examples/base_server/server_async_secure.py +0 -0
  40. {digitalkin-0.2.13 → digitalkin-0.2.15}/examples/base_server/server_sync_insecure.py +0 -0
  41. {digitalkin-0.2.13 → digitalkin-0.2.15}/examples/base_server/server_sync_secure.py +0 -0
  42. {digitalkin-0.2.13 → digitalkin-0.2.15}/examples/modules/__init__.py +0 -0
  43. {digitalkin-0.2.13 → digitalkin-0.2.15}/setup.cfg +0 -0
  44. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/__init__.py +0 -0
  45. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/grpc_servers/__init__.py +0 -0
  46. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/grpc_servers/_base_server.py +0 -0
  47. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/grpc_servers/module_server.py +0 -0
  48. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/grpc_servers/registry_server.py +0 -0
  49. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/grpc_servers/registry_servicer.py +0 -0
  50. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/grpc_servers/utils/factory.py +0 -0
  51. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/grpc_servers/utils/grpc_client_wrapper.py +0 -0
  52. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/grpc_servers/utils/models.py +0 -0
  53. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/grpc_servers/utils/types.py +0 -0
  54. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/logger.py +0 -0
  55. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/models/__init__.py +0 -0
  56. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/models/module/module.py +0 -0
  57. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/models/services/__init__.py +0 -0
  58. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/models/services/cost.py +0 -0
  59. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/models/services/storage.py +0 -0
  60. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/modules/__init__.py +0 -0
  61. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/modules/job_manager/job_manager_models.py +0 -0
  62. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/py.typed +0 -0
  63. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/services/__init__.py +0 -0
  64. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/services/agent/__init__.py +0 -0
  65. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/services/agent/agent_strategy.py +0 -0
  66. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/services/agent/default_agent.py +0 -0
  67. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/services/base_strategy.py +0 -0
  68. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/services/cost/__init__.py +0 -0
  69. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/services/cost/cost_strategy.py +0 -0
  70. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/services/cost/default_cost.py +0 -0
  71. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/services/cost/grpc_cost.py +0 -0
  72. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/services/filesystem/__init__.py +0 -0
  73. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/services/identity/__init__.py +0 -0
  74. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/services/identity/default_identity.py +0 -0
  75. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/services/identity/identity_strategy.py +0 -0
  76. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/services/registry/__init__.py +0 -0
  77. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/services/registry/default_registry.py +0 -0
  78. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/services/registry/registry_strategy.py +0 -0
  79. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/services/services_config.py +0 -0
  80. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/services/services_models.py +0 -0
  81. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/services/setup/__init__.py +0 -0
  82. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/services/setup/default_setup.py +0 -0
  83. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/services/setup/setup_strategy.py +0 -0
  84. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/services/snapshot/__init__.py +0 -0
  85. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/services/snapshot/default_snapshot.py +0 -0
  86. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/services/snapshot/snapshot_strategy.py +0 -0
  87. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/services/storage/__init__.py +0 -0
  88. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/services/storage/default_storage.py +0 -0
  89. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/services/storage/grpc_storage.py +0 -0
  90. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/services/storage/storage_strategy.py +0 -0
  91. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/utils/__init__.py +0 -0
  92. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/utils/arg_parser.py +0 -0
  93. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/utils/development_mode_action.py +0 -0
  94. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin/utils/llm_ready_schema.py +0 -0
  95. {digitalkin-0.2.13 → digitalkin-0.2.15}/src/digitalkin.egg-info/dependency_links.txt +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: digitalkin
3
- Version: 0.2.13
3
+ Version: 0.2.15
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,36 +452,36 @@ 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>=0.1.10
455
+ Requires-Dist: digitalkin-proto>=0.1.15
456
456
  Requires-Dist: grpcio-health-checking>=1.71.0
457
457
  Requires-Dist: grpcio-reflection>=1.71.0
458
458
  Requires-Dist: grpcio-status>=1.71.0
459
- Requires-Dist: pydantic>=2.11.4
459
+ Requires-Dist: pydantic>=2.11.5
460
460
  Provides-Extra: dev
461
- Requires-Dist: typos>=1.32.0; extra == "dev"
462
- Requires-Dist: ruff>=0.11.9; extra == "dev"
463
- Requires-Dist: mypy>=1.15.0; extra == "dev"
464
- Requires-Dist: pyright>=1.1.400; extra == "dev"
461
+ Requires-Dist: typos>=1.33.1; extra == "dev"
462
+ Requires-Dist: ruff>=0.11.13; extra == "dev"
463
+ Requires-Dist: mypy>=1.16.0; extra == "dev"
464
+ Requires-Dist: pyright>=1.1.401; extra == "dev"
465
465
  Requires-Dist: pre-commit>=4.2.0; extra == "dev"
466
466
  Requires-Dist: bump2version>=1.0.1; extra == "dev"
467
467
  Requires-Dist: build>=1.2.2; extra == "dev"
468
468
  Requires-Dist: twine>=6.1.0; extra == "dev"
469
- Requires-Dist: cryptography>=44.0.3; extra == "dev"
470
- Requires-Dist: taskiq[reload]>=0.11.17; extra == "dev"
469
+ Requires-Dist: cryptography>=45.0.4; extra == "dev"
471
470
  Provides-Extra: examples
472
471
  Requires-Dist: openai>=1.75.0; extra == "examples"
473
472
  Provides-Extra: tests
474
- Requires-Dist: freezegun>=1.5.1; extra == "tests"
473
+ Requires-Dist: freezegun>=1.5.2; extra == "tests"
475
474
  Requires-Dist: hdrhistogram>=0.10.3; extra == "tests"
476
475
  Requires-Dist: grpcio-testing>=1.71.0; extra == "tests"
477
476
  Requires-Dist: psutil>=7.0.0; extra == "tests"
478
- Requires-Dist: pytest>=8.3.4; extra == "tests"
479
- Requires-Dist: pytest-asyncio>=0.26.0; extra == "tests"
477
+ Requires-Dist: pytest>=8.4.0; extra == "tests"
478
+ Requires-Dist: pytest-asyncio>=1.0.0; extra == "tests"
480
479
  Requires-Dist: pytest-cov>=6.1.0; extra == "tests"
481
480
  Provides-Extra: taskiq
482
- Requires-Dist: rstream>=0.20.9; extra == "taskiq"
481
+ Requires-Dist: rstream>=0.30.0; extra == "taskiq"
483
482
  Requires-Dist: taskiq-aio-pika>=0.4.2; extra == "taskiq"
484
- Requires-Dist: taskiq-redis>=1.0.8; extra == "taskiq"
483
+ Requires-Dist: taskiq-redis>=1.0.9; extra == "taskiq"
484
+ Requires-Dist: taskiq[reload]>=0.11.17; extra == "taskiq"
485
485
  Dynamic: license-file
486
486
 
487
487
  # DigitalKin Python SDK
@@ -127,6 +127,16 @@ class CPUOutput(BaseModel):
127
127
  )
128
128
 
129
129
 
130
+ class CPUConfigSetup(BaseModel):
131
+ """Config Setup model definining data that will be pre-computed for each setup and module instance."""
132
+
133
+ files: list[str] = Field(
134
+ ...,
135
+ title="Files to embed",
136
+ description="List of files to embed in the setup lifecycle.",
137
+ )
138
+
139
+
130
140
  class CPUSetup(BaseModel):
131
141
  """Setup model defining module configuration parameters."""
132
142
 
@@ -175,7 +185,7 @@ client_config = ClientConfig(
175
185
  )
176
186
 
177
187
 
178
- class CPUIntensiveModule(BaseModule[CPUInput, CPUOutput, CPUSetup, CPUToolSecret]):
188
+ class CPUIntensiveModule(BaseModule[CPUInput, CPUOutput, CPUSetup, CPUToolSecret, None]):
179
189
  """A CPU endpoint tool module module."""
180
190
 
181
191
  name = "CPUIntensiveModule"
@@ -202,7 +212,6 @@ class CPUIntensiveModule(BaseModule[CPUInput, CPUOutput, CPUSetup, CPUToolSecret
202
212
  "client_config": client_config,
203
213
  },
204
214
  "filesystem": {
205
- "config": {},
206
215
  "client_config": client_config,
207
216
  },
208
217
  "cost": {
@@ -11,7 +11,6 @@ from pydantic import BaseModel, Field
11
11
  from digitalkin.grpc_servers.utils.models import ClientConfig, SecurityMode, ServerMode
12
12
  from digitalkin.modules._base_module import BaseModule
13
13
  from digitalkin.services.services_models import ServicesStrategy
14
- from digitalkin.services.setup.setup_strategy import SetupData
15
14
 
16
15
  # Configure logging with clear formatting
17
16
  logging.basicConfig(
@@ -154,6 +153,16 @@ class OpenAISetup(BaseModel):
154
153
  )
155
154
 
156
155
 
156
+ class OpenAIConfigSetup(BaseModel):
157
+ """Setup model defining module configuration parameters."""
158
+
159
+ rag_files: list[bytes] = Field(
160
+ ...,
161
+ title="RAG Files",
162
+ description="Files used for retrieval-augmented generation (RAG) with the OpenAI module.",
163
+ )
164
+
165
+
157
166
  class OpenAIToolSecret(BaseModel):
158
167
  """Secret model defining module configuration parameters."""
159
168
 
@@ -167,13 +176,22 @@ client_config = ClientConfig(
167
176
  )
168
177
 
169
178
 
170
- class OpenAIToolModule(BaseModule[OpenAIInput, OpenAIOutput, OpenAISetup, OpenAIToolSecret]):
179
+ class OpenAIToolModule(
180
+ BaseModule[
181
+ OpenAIInput,
182
+ OpenAIOutput,
183
+ OpenAISetup,
184
+ OpenAIToolSecret,
185
+ OpenAIConfigSetup,
186
+ ]
187
+ ):
171
188
  """A openAI endpoint tool module module."""
172
189
 
173
190
  name = "OpenAIToolModule"
174
191
  description = "A module that interacts with OpenAI API to process text"
175
192
 
176
193
  # Define the schema formats for the module
194
+ config_setup_format = OpenAIConfigSetup
177
195
  input_format = OpenAIInput
178
196
  output_format = OpenAIOutput
179
197
  setup_format = OpenAISetup
@@ -205,7 +223,27 @@ class OpenAIToolModule(BaseModule[OpenAIInput, OpenAIOutput, OpenAISetup, OpenAI
205
223
  },
206
224
  }
207
225
 
208
- async def initialize(self, setup_data: SetupData) -> None:
226
+ async def run_config_setup(
227
+ self,
228
+ config_setup_data: OpenAIConfigSetup,
229
+ setup_data: OpenAISetup,
230
+ callback: Callable,
231
+ ) -> None:
232
+ """Configure the module with additional setup data.
233
+
234
+ Args:
235
+ config_setup_data: Additional configuration content.
236
+ setup_data: Initial setup data for the module.
237
+ callback: Function to send output data back to the client.
238
+ """
239
+ logger.info("Configuring OpenAIToolModule with additional setup data. %s", config_setup_data)
240
+
241
+ # Here you can process config_content and update setup_data as needed
242
+ # For now, we just return the original setup_data
243
+ setup_data.developer_prompt = "| + |".join(f.decode("utf-8") for f in config_setup_data.rag_files)
244
+ await callback(setup_data)
245
+
246
+ async def initialize(self, setup_data: OpenAISetup) -> None:
209
247
  """Initialize the module capabilities.
210
248
 
211
249
  This method is called when the module is loaded by the server.
@@ -93,7 +93,6 @@ class TextTransformModule(BaseModule[TextTransformInput, TextTransformOutput, Te
93
93
  "client_config": client_config,
94
94
  },
95
95
  "filesystem": {
96
- "config": {},
97
96
  "client_config": client_config,
98
97
  },
99
98
  }
@@ -0,0 +1,198 @@
1
+ """Example module implementation to test ArchetypeModule functionality."""
2
+
3
+ import asyncio
4
+ import datetime
5
+ from collections.abc import Callable
6
+ from typing import TYPE_CHECKING, Any
7
+
8
+ from pydantic import BaseModel, Field
9
+
10
+ from digitalkin.logger import logger
11
+ from digitalkin.models.module import ModuleStatus
12
+ from digitalkin.modules.archetype_module import ArchetypeModule
13
+ from digitalkin.services.services_config import ServicesConfig
14
+ from digitalkin.services.services_models import ServicesMode
15
+ from digitalkin.services.filesystem.filesystem_strategy import FilesystemRecord, FileFilter, UploadFileData
16
+
17
+
18
+ class ExampleInput(BaseModel):
19
+ """Input model for example module."""
20
+
21
+ message: str = Field(description="Message to process")
22
+ number: int = Field(description="Number to process")
23
+
24
+
25
+ class ExampleOutput(BaseModel):
26
+ """Output model for example module."""
27
+
28
+ processed_message: str = Field(description="The processed message")
29
+ processed_number: int = Field(description="The processed number")
30
+ timestamp: datetime.datetime = Field(description="When the processing was done")
31
+
32
+
33
+ class ExampleSetup(BaseModel):
34
+ """Setup model for example module."""
35
+
36
+ processing_mode: str = Field(description="Mode to process data in", default="default")
37
+ multiply_factor: int = Field(description="Factor to multiply number by", default=1)
38
+
39
+
40
+ class ExampleSecret(BaseModel):
41
+ """Secret model for example module."""
42
+
43
+ api_key: str = Field(description="API key for external service")
44
+
45
+
46
+ class ExampleStorage(BaseModel):
47
+ """Secret model for example module."""
48
+
49
+ test_key: str = Field(description="Test value for storage")
50
+
51
+
52
+ class ExampleModule(ArchetypeModule[ExampleInput, ExampleOutput, ExampleSetup, ExampleSecret, None]):
53
+ """Example module that demonstrates ArchetypeModule functionality."""
54
+
55
+ name = "ExampleModule"
56
+ description = "An example module for testing purposes"
57
+ input_format = ExampleInput
58
+ output_format = ExampleOutput
59
+ setup_format = ExampleSetup
60
+ secret_format = ExampleSecret
61
+ metadata = {"name": "ExampleModule", "description": "A module for testing ArchetypeModule functionality"}
62
+
63
+ # Define services_config_params with default values
64
+ services_config_strategies = {}
65
+ services_config_params = {"cost": {"config": {}}, "storage": {"config": {}}} # Filesystem has no config but it's enabled
66
+
67
+ def __init__(self, job_id: str, mission_id: str, setup_version_id: str) -> None:
68
+ """Initialize the example module.
69
+
70
+ Args:
71
+ job_id: Unique identifier for the job
72
+ name: Optional name for the module
73
+ """
74
+ # Initialize services configuration using the class attribute before the instance is created
75
+ self.services_config = ServicesConfig(
76
+ services_config_strategies=self.services_config_strategies,
77
+ services_config_params=self.services_config_params,
78
+ mode=ServicesMode.LOCAL,
79
+ )
80
+
81
+ super().__init__(job_id, mission_id, setup_version_id)
82
+
83
+ async def initialize(self, setup_data: ExampleSetup) -> None:
84
+ """Initialize the module.
85
+
86
+ Args:
87
+ setup_data: Setup data for the module
88
+ """
89
+ logger.info("Initializing ExampleModule with setup data: %s", setup_data)
90
+ self.setup = self.setup_format.model_validate(setup_data)
91
+ logger.info("Initialization complete, using processing mode: [%s]", self.setup.processing_mode)
92
+
93
+ async def run(
94
+ self,
95
+ input_data: dict[str, Any],
96
+ setup_data: ExampleSetup,
97
+ callback: Callable,
98
+ ) -> None:
99
+ """Run the module.
100
+
101
+ Args:
102
+ input_data: Input data for the module
103
+ setup_data: Setup data for the module
104
+ callback: Callback function to report progress
105
+ """
106
+ # Validate the input data
107
+ input_model = self.input_format.model_validate(input_data)
108
+ logger.info("Running with input data: %s", input_model)
109
+
110
+ # Process the data
111
+ processed_message = f"Processed: {input_model.message}"
112
+ processed_number = input_model.number * self.setup.multiply_factor
113
+
114
+ # Create output model
115
+ file = UploadFileData(
116
+ content=b"%s\n%s" % (processed_message.encode(), str(processed_number).encode()),
117
+ name="example_output.txt",
118
+ file_type="text/plain",
119
+ content_type="text/plain",
120
+ metadata={"example_key": "example_value"},
121
+ replace_if_exists=True,
122
+ )
123
+
124
+ records, uploaded, failed = self.filesystem.upload_files(files=[file])
125
+ for record in records:
126
+ logger.info("Uploaded file: %s, uploaded: %d, failed: %d", record, uploaded, failed)
127
+ logger.info("Stored file with ID: %s", record.id)
128
+ callback(record.model_dump())
129
+ # Call the callback with the output data
130
+
131
+
132
+ # Wait a bit to simulate processing time
133
+ await asyncio.sleep(1)
134
+
135
+ async def run_config_setup(
136
+ self,
137
+ setup_data: ExampleSetup,
138
+ ) -> None:
139
+ """Run the configuration setup for the module.
140
+
141
+ Args:
142
+ setup_data: Setup data for the module
143
+ """
144
+ logger.info("Running config setup with data: %s", setup_data)
145
+ # Here we could implement any additional configuration logic if needed
146
+
147
+ async def cleanup(self) -> None:
148
+ """Clean up the module."""
149
+ logger.info("Cleaning up ExampleModule")
150
+ # Nothing to clean up in this example
151
+
152
+
153
+ async def test_module() -> None:
154
+ """Test the example module."""
155
+ # Create the module
156
+ module = ExampleModule(job_id="test-job-123", mission_id="test-mission-123", setup_version_id="test-setup-123")
157
+
158
+ # Define input and setup data
159
+ input_data = ExampleInput(message="Hello, world!", number=42)
160
+
161
+ setup_data = ExampleSetup(processing_mode="test", multiply_factor=10)
162
+
163
+ # Define a callback function
164
+ def callback(result) -> None:
165
+ logger.info(f"callback {result}")
166
+
167
+ # Start the module
168
+ await module.start(input_data, setup_data, callback)
169
+
170
+ # Wait for the module to complete
171
+ while module.status not in {ModuleStatus.STOPPED, ModuleStatus.FAILED}:
172
+ await asyncio.sleep(0.5)
173
+
174
+ # Check the storage
175
+ if module.status == ModuleStatus.STOPPED:
176
+ files, nb_results = module.filesystem.get_files(
177
+ filters=FileFilter(name="example_output.txt", context="test-mission-123"),
178
+ )
179
+ for file in files:
180
+ module.filesystem.update_file(file.id, file_type="updated")
181
+ # module.filesystem.delete_files(filters=FileFilter(name="example_output.txt", context="test-mission-123"), permanent=True)
182
+
183
+ logger.info("Retrieved file: %s with ID: %s", file.name, file.id)
184
+ try:
185
+ file_record = module.filesystem.get_file(file_id=file.id, include_content=True)
186
+ if file_record:
187
+ logger.info("File ID: %s", file_record.id)
188
+ logger.info("File name: %s", file_record.name)
189
+ logger.info("File type: %s", file_record.file_type)
190
+ logger.info("File status: %s", file_record.status)
191
+ logger.info("File content: %s", file_record.content.decode())
192
+ except Exception:
193
+ logger.error("No file retrieved")
194
+
195
+
196
+ if __name__ == "__main__":
197
+ # Run the module test
198
+ asyncio.run(test_module())
@@ -51,7 +51,7 @@ class ExampleStorage(BaseModel):
51
51
  test_key: str = Field(description="Test value for storage")
52
52
 
53
53
 
54
- class ExampleModule(ArchetypeModule[ExampleInput, ExampleOutput, ExampleSetup, ExampleSecret]):
54
+ class ExampleModule(ArchetypeModule[ExampleInput, ExampleOutput, ExampleSetup, ExampleSecret, None]):
55
55
  """Example module that demonstrates ArchetypeModule functionality."""
56
56
 
57
57
  name = "ExampleModule"
@@ -64,14 +64,15 @@ class ExampleModule(ArchetypeModule[ExampleInput, ExampleOutput, ExampleSetup, E
64
64
 
65
65
  # Define services_config_params with default values
66
66
  services_config_strategies = {}
67
- services_config_params = {"storage": {"config": {"example": ExampleOutput}}, "filesystem": {"config": {}}}
67
+ services_config_params = {"storage": {"config": {"example": ExampleOutput}},"cost": {"config":{}}}
68
68
 
69
- def __init__(self, job_id: str, mission_id: str) -> None:
69
+ def __init__(self, job_id: str, mission_id: str, setup_version_id: str) -> None:
70
70
  """Initialize the example module.
71
71
 
72
72
  Args:
73
73
  job_id: Unique identifier for the job
74
- name: Optional name for the module
74
+ mission_id: Unique identifier for the mission
75
+ setup_version_id: Unique identifier for the setup version
75
76
  """
76
77
  # Initialize services configuration using the class attribute before the instance is created
77
78
  self.services_config = ServicesConfig(
@@ -80,7 +81,7 @@ class ExampleModule(ArchetypeModule[ExampleInput, ExampleOutput, ExampleSetup, E
80
81
  mode=ServicesMode.LOCAL,
81
82
  )
82
83
 
83
- super().__init__(job_id, mission_id)
84
+ super().__init__(job_id, mission_id, setup_version_id)
84
85
 
85
86
  async def initialize(self, setup_data: ExampleSetup) -> None:
86
87
  """Initialize the module.
@@ -92,6 +93,18 @@ class ExampleModule(ArchetypeModule[ExampleInput, ExampleOutput, ExampleSetup, E
92
93
  self.setup = self.setup_format.model_validate(setup_data)
93
94
  logger.info("Initialization complete, using processing mode: [%s]", self.setup.processing_mode)
94
95
 
96
+ async def run_config_setup(
97
+ self,
98
+ setup_data: ExampleSetup,
99
+ ) -> None:
100
+ """Run the configuration setup for the module.
101
+
102
+ Args:
103
+ setup_data: Setup data for the module
104
+ """
105
+ logger.info("Running config setup with data: %s", setup_data)
106
+ # Here we could implement any additional configuration logic if needed
107
+
95
108
  async def run(
96
109
  self,
97
110
  input_data: dict[str, Any],
@@ -142,7 +155,7 @@ class ExampleModule(ArchetypeModule[ExampleInput, ExampleOutput, ExampleSetup, E
142
155
  async def test_module() -> None:
143
156
  """Test the example module."""
144
157
  # Create the module
145
- module = ExampleModule(job_id="test-job-123", mission_id="test-mission-123")
158
+ module = ExampleModule(job_id="test-job-123", mission_id="test-mission-123", setup_version_id="test-setup-123")
146
159
 
147
160
  # Define input and setup data
148
161
  input_data = ExampleInput(message="Hello, world!", number=42)
@@ -171,7 +184,7 @@ async def test_module() -> None:
171
184
  def test_storage_directly() -> None:
172
185
  """Test the storage service directly."""
173
186
  # Initialize storage service
174
- storage = ServicesConfig().storage(mission_id="test-mission", config={"example": ExampleStorage})
187
+ storage = ServicesConfig().storage(mission_id="test-mission",setup_version_id="test-setup-123", config={"example": ExampleStorage})
175
188
 
176
189
  # Create a test record
177
190
  storage.store("example", "test_table", {"test_key": "test_value"}, "OUTPUT")
@@ -12,7 +12,7 @@
12
12
 
13
13
  keywords = [ "digitalkin", "kin", "agent", "gprc", "sdk" ]
14
14
  # Version of the package automatically updated by bump2version (that is why it is separated)
15
- version = "0.2.13"
15
+ version = "0.2.15"
16
16
 
17
17
  classifiers = [
18
18
  "Development Status :: 3 - Alpha",
@@ -29,41 +29,41 @@
29
29
  ]
30
30
 
31
31
  dependencies = [
32
- "digitalkin-proto>=0.1.10",
32
+ "digitalkin-proto>=0.1.15",
33
33
  "grpcio-health-checking>=1.71.0",
34
34
  "grpcio-reflection>=1.71.0",
35
35
  "grpcio-status>=1.71.0",
36
- "pydantic>=2.11.4",
36
+ "pydantic>=2.11.5",
37
37
  ]
38
38
 
39
39
  [project.optional-dependencies]
40
40
  dev = [
41
- "typos>=1.32.0",
42
- "ruff>=0.11.9",
43
- "mypy>=1.15.0",
44
- "pyright>=1.1.400",
41
+ "typos>=1.33.1",
42
+ "ruff>=0.11.13",
43
+ "mypy>=1.16.0",
44
+ "pyright>=1.1.401",
45
45
  "pre-commit>=4.2.0",
46
46
  "bump2version>=1.0.1",
47
47
  "build>=1.2.2",
48
48
  "twine>=6.1.0",
49
- "cryptography>=44.0.3",
50
- "taskiq[reload]>=0.11.17",
49
+ "cryptography>=45.0.4",
51
50
  ]
52
51
  examples = [ "openai>=1.75.0" ]
53
52
  tests = [
54
- "freezegun>=1.5.1",
53
+ "freezegun>=1.5.2",
55
54
  "hdrhistogram>=0.10.3",
56
55
  "grpcio-testing>=1.71.0",
57
56
  "psutil>=7.0.0",
58
- "pytest>=8.3.4",
59
- "pytest-asyncio>=0.26.0",
57
+ "pytest>=8.4.0",
58
+ "pytest-asyncio>=1.0.0",
60
59
  "pytest-cov>=6.1.0",
61
60
  ]
62
61
  taskiq = [
63
- "rstream>=0.20.9",
62
+ "rstream>=0.30.0",
64
63
  "taskiq-aio-pika>=0.4.2",
65
- "taskiq-redis>=1.0.8",
66
- ]
64
+ "taskiq-redis>=1.0.9",
65
+ "taskiq[reload]>=0.11.17",
66
+ ]
67
67
 
68
68
  [project.urls]
69
69
  Homepage = "https://github.com/DigitalKin-ai/digitalkin"
@@ -115,8 +115,8 @@
115
115
  "tool_module.py",
116
116
  "trigger_module.py",
117
117
  "generate_certificates.py",
118
- "test_*.py",
119
- "tests/*",
118
+ # "test_*.py",
119
+ # "tests/*",
120
120
  "examples/*",
121
121
  ]
122
122
  preview = true
@@ -186,9 +186,33 @@
186
186
  "TRY002",
187
187
  "FBT001",
188
188
  "FBT002",
189
+ "ANN001",
189
190
  "ANN003",
190
191
  "ANN002",
191
192
  "ARG002",
193
+ "PLR0914",
194
+ "PLR2004",
195
+ "PLR0915",
196
+ "PLR1702",
197
+ "D417",
198
+ "ARG001",
199
+ "D103",
200
+ "D102",
201
+ "D101",
202
+ "D107",
203
+ "N802",
204
+ "DOC201",
205
+ "DOC501",
206
+ "PERF203",
207
+ "C901",
208
+ "PLR0911",
209
+ "G004",
210
+ "ERA001",
211
+ "E501",
212
+ "ANN401",
213
+ "PLR0912",
214
+ "ANN201",
215
+ "D100"
192
216
  ]
193
217
 
194
218
  [tool.ruff.lint.pylint]
@@ -5,4 +5,4 @@ from importlib.metadata import PackageNotFoundError, version
5
5
  try:
6
6
  __version__ = version("digitalkin")
7
7
  except PackageNotFoundError:
8
- __version__ = "0.2.13"
8
+ __version__ = "0.2.15"
@@ -81,6 +81,62 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
81
81
  )
82
82
  self.setup = GrpcSetup() if self.args.services_mode == ServicesMode.REMOTE else DefaultSetup()
83
83
 
84
+ async def ConfigSetupModule( # noqa: N802
85
+ self,
86
+ request: lifecycle_pb2.ConfigSetupModuleRequest,
87
+ context: grpc.aio.ServicerContext,
88
+ ) -> lifecycle_pb2.ConfigSetupModuleResponse:
89
+ """Configure the module setup.
90
+
91
+ Args:
92
+ request: The configuration request.
93
+ context: The gRPC context.
94
+
95
+ Returns:
96
+ A response indicating success or failure.
97
+
98
+ Raises:
99
+ ServicerError: if the setup data is not returned or job creation fails.
100
+ """
101
+ logger.info("ConfigSetupVersion called for module: '%s'", self.module_class.__name__)
102
+ # Process the module input
103
+ # TODO: Secret should be used here as well
104
+ setup_version = request.setup_version
105
+ config_setup_data = self.module_class.create_config_setup_model(json_format.MessageToDict(request.content))
106
+ setup_version_data = self.module_class.create_setup_model(
107
+ json_format.MessageToDict(request.setup_version.content)
108
+ )
109
+
110
+ if not setup_version_data:
111
+ msg = "No setup data returned."
112
+ raise ServicerError(msg)
113
+
114
+ if not config_setup_data:
115
+ msg = "No config setup data returned."
116
+ raise ServicerError(msg)
117
+
118
+ # create a task to run the module in background
119
+ job_id = await self.job_manager.create_config_setup_instance_job(
120
+ config_setup_data,
121
+ setup_version_data,
122
+ request.mission_id,
123
+ setup_version.id,
124
+ )
125
+
126
+ if job_id is None:
127
+ context.set_code(grpc.StatusCode.NOT_FOUND)
128
+ context.set_details("Failed to create module instance")
129
+ return lifecycle_pb2.ConfigSetupModuleResponse(success=False)
130
+
131
+ updated_setup_data = await self.job_manager.generate_config_setup_module_response(job_id)
132
+ logger.warning(f"Updated setup data: {updated_setup_data=}")
133
+ setup_version.content = json_format.ParseDict(
134
+ updated_setup_data,
135
+ struct_pb2.Struct(),
136
+ ignore_unknown_fields=True,
137
+ )
138
+ return lifecycle_pb2.ConfigSetupModuleResponse(success=True, setup_version=setup_version)
139
+
84
140
  async def StartModule( # noqa: N802
85
141
  self,
86
142
  request: lifecycle_pb2.StartModuleRequest,
@@ -116,7 +172,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
116
172
  setup_data = self.module_class.create_setup_model(setup_data_class.current_setup_version.content)
117
173
 
118
174
  # create a task to run the module in background
119
- job_id = await self.job_manager.create_job(
175
+ job_id = await self.job_manager.create_module_instance_job(
120
176
  input_data,
121
177
  setup_data,
122
178
  mission_id=request.mission_id,
@@ -391,3 +447,39 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
391
447
  success=True,
392
448
  secret_schema=secret_format_struct,
393
449
  )
450
+
451
+ async def GetConfigSetupModule( # noqa: N802
452
+ self,
453
+ request: information_pb2.GetConfigSetupModuleRequest,
454
+ context: grpc.ServicerContext,
455
+ ) -> information_pb2.GetConfigSetupModuleResponse:
456
+ """Get information about the module's setup and configuration.
457
+
458
+ Args:
459
+ request: The get module setup request.
460
+ context: The gRPC context.
461
+
462
+ Returns:
463
+ A response with the module's setup information.
464
+ """
465
+ logger.debug("GetConfigSetupModule called for module: '%s'", self.module_class.__name__)
466
+
467
+ # Get setup schema if available
468
+ try:
469
+ # Convert schema to proto format
470
+ config_setup_schema_proto = self.module_class.get_config_setup_format(llm_format=request.llm_format)
471
+ config_setup_format_struct = json_format.Parse(
472
+ text=config_setup_schema_proto,
473
+ message=struct_pb2.Struct(), # pylint: disable=no-member
474
+ ignore_unknown_fields=True,
475
+ )
476
+ except NotImplementedError as e:
477
+ logger.warning(e)
478
+ context.set_code(grpc.StatusCode.UNIMPLEMENTED)
479
+ context.set_details(e)
480
+ return information_pb2.GetConfigSetupModuleResponse()
481
+
482
+ return information_pb2.GetConfigSetupModuleResponse(
483
+ success=True,
484
+ config_setup_schema=config_setup_format_struct,
485
+ )