digitalkin 0.2.6__tar.gz → 0.2.8__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.2.6 → digitalkin-0.2.8}/PKG-INFO +7 -7
- {digitalkin-0.2.6 → digitalkin-0.2.8}/examples/modules/storage_module.py +9 -4
- {digitalkin-0.2.6 → digitalkin-0.2.8}/pyproject.toml +7 -7
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/__version__.py +1 -1
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/grpc_servers/module_servicer.py +39 -3
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/modules/_base_module.py +9 -8
- digitalkin-0.2.8/src/digitalkin/services/storage/default_storage.py +220 -0
- digitalkin-0.2.8/src/digitalkin/services/storage/grpc_storage.py +197 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/storage/storage_strategy.py +156 -104
- digitalkin-0.2.8/src/digitalkin/utils/llm_ready_schema.py +75 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin.egg-info/PKG-INFO +7 -7
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin.egg-info/SOURCES.txt +2 -1
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin.egg-info/requires.txt +6 -6
- digitalkin-0.2.6/src/digitalkin/services/storage/default_storage.py +0 -218
- digitalkin-0.2.6/src/digitalkin/services/storage/grpc_storage.py +0 -156
- {digitalkin-0.2.6 → digitalkin-0.2.8}/LICENSE +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/README.md +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/examples/base_server/__init__.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/examples/base_server/mock/__init__.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/examples/base_server/mock/mock_pb2.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/examples/base_server/mock/mock_pb2_grpc.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/examples/base_server/server_async_insecure.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/examples/base_server/server_async_secure.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/examples/base_server/server_sync_insecure.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/examples/base_server/server_sync_secure.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/examples/modules/__init__.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/examples/modules/minimal_llm_module.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/examples/modules/text_transform_module.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/setup.cfg +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/__init__.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/grpc_servers/__init__.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/grpc_servers/_base_server.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/grpc_servers/module_server.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/grpc_servers/registry_server.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/grpc_servers/registry_servicer.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/grpc_servers/utils/exceptions.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/grpc_servers/utils/factory.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/grpc_servers/utils/grpc_client_wrapper.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/grpc_servers/utils/models.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/grpc_servers/utils/types.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/logger.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/models/__init__.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/models/module/__init__.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/models/module/module.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/models/module/module_types.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/models/services/__init__.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/models/services/cost.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/models/services/storage.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/modules/__init__.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/modules/archetype_module.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/modules/job_manager.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/modules/tool_module.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/modules/trigger_module.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/py.typed +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/__init__.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/agent/__init__.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/agent/agent_strategy.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/agent/default_agent.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/base_strategy.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/cost/__init__.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/cost/cost_strategy.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/cost/default_cost.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/cost/grpc_cost.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/filesystem/__init__.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/filesystem/default_filesystem.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/filesystem/filesystem_strategy.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/filesystem/grpc_filesystem.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/identity/__init__.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/identity/default_identity.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/identity/identity_strategy.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/registry/__init__.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/registry/default_registry.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/registry/registry_strategy.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/services_config.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/services_models.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/setup/default_setup.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/setup/grpc_setup.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/setup/setup_strategy.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/snapshot/__init__.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/snapshot/default_snapshot.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/snapshot/snapshot_strategy.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/storage/__init__.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/utils/__init__.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/utils/arg_parser.py +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin.egg-info/dependency_links.txt +0 -0
- {digitalkin-0.2.6 → digitalkin-0.2.8}/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.2.
|
|
3
|
+
Version: 0.2.8
|
|
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,20 +452,20 @@ 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.
|
|
455
|
+
Requires-Dist: digitalkin-proto>=0.1.8
|
|
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: openai>=1.
|
|
460
|
-
Requires-Dist: pydantic>=2.11.
|
|
459
|
+
Requires-Dist: openai>=1.76.2
|
|
460
|
+
Requires-Dist: pydantic>=2.11.4
|
|
461
461
|
Provides-Extra: dev
|
|
462
462
|
Requires-Dist: pytest>=8.3.4; extra == "dev"
|
|
463
463
|
Requires-Dist: pytest-asyncio>=0.26.0; extra == "dev"
|
|
464
464
|
Requires-Dist: pytest-cov>=6.1.0; extra == "dev"
|
|
465
|
-
Requires-Dist: typos>=1.31.
|
|
466
|
-
Requires-Dist: ruff>=0.11.
|
|
465
|
+
Requires-Dist: typos>=1.31.2; extra == "dev"
|
|
466
|
+
Requires-Dist: ruff>=0.11.7; extra == "dev"
|
|
467
467
|
Requires-Dist: mypy>=1.15.0; extra == "dev"
|
|
468
|
-
Requires-Dist: pyright>=1.1.
|
|
468
|
+
Requires-Dist: pyright>=1.1.400; extra == "dev"
|
|
469
469
|
Requires-Dist: pre-commit>=4.2.0; extra == "dev"
|
|
470
470
|
Requires-Dist: bump2version>=1.0.1; extra == "dev"
|
|
471
471
|
Requires-Dist: build>=1.2.2; extra == "dev"
|
|
@@ -119,7 +119,12 @@ class ExampleModule(ArchetypeModule[ExampleInput, ExampleOutput, ExampleSetup, E
|
|
|
119
119
|
)
|
|
120
120
|
|
|
121
121
|
# Store the output data in storage
|
|
122
|
-
storage_id = self.storage.store(
|
|
122
|
+
storage_id = self.storage.store(
|
|
123
|
+
collection="example",
|
|
124
|
+
record_id=f"example_outputs",
|
|
125
|
+
data=output_data.model_dump(),
|
|
126
|
+
data_type="OUTPUT"
|
|
127
|
+
)
|
|
123
128
|
|
|
124
129
|
logger.info("Stored output data with ID: %s", storage_id)
|
|
125
130
|
|
|
@@ -159,7 +164,7 @@ async def test_module() -> None:
|
|
|
159
164
|
|
|
160
165
|
# Check the storage
|
|
161
166
|
if module.status == ModuleStatus.STOPPED:
|
|
162
|
-
result: StorageRecord = module.storage.read("example_outputs")
|
|
167
|
+
result: StorageRecord = module.storage.read("example", "example_outputs")
|
|
163
168
|
if result:
|
|
164
169
|
pass
|
|
165
170
|
|
|
@@ -170,10 +175,10 @@ def test_storage_directly() -> None:
|
|
|
170
175
|
storage = ServicesConfig().storage(mission_id="test-mission", config={"test_table": ExampleStorage})
|
|
171
176
|
|
|
172
177
|
# Create a test record
|
|
173
|
-
storage.store("test_table", {"test_key": "test_value"}, "OUTPUT")
|
|
178
|
+
storage.store("example", "test_table", {"test_key": "test_value"}, "OUTPUT")
|
|
174
179
|
|
|
175
180
|
# Retrieve the record
|
|
176
|
-
retrieved = storage.read("test_table")
|
|
181
|
+
retrieved = storage.read("example", "test_table")
|
|
177
182
|
|
|
178
183
|
if retrieved:
|
|
179
184
|
pass
|
|
@@ -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.
|
|
15
|
+
version = "0.2.8"
|
|
16
16
|
|
|
17
17
|
classifiers = [
|
|
18
18
|
"Development Status :: 3 - Alpha",
|
|
@@ -29,12 +29,12 @@
|
|
|
29
29
|
]
|
|
30
30
|
|
|
31
31
|
dependencies = [
|
|
32
|
-
"digitalkin-proto>=0.1.
|
|
32
|
+
"digitalkin-proto>=0.1.8",
|
|
33
33
|
"grpcio-health-checking>=1.71.0",
|
|
34
34
|
"grpcio-reflection>=1.71.0",
|
|
35
35
|
"grpcio-status>=1.71.0",
|
|
36
|
-
"openai>=1.
|
|
37
|
-
"pydantic>=2.11.
|
|
36
|
+
"openai>=1.76.2",
|
|
37
|
+
"pydantic>=2.11.4",
|
|
38
38
|
]
|
|
39
39
|
|
|
40
40
|
[project.optional-dependencies]
|
|
@@ -42,10 +42,10 @@
|
|
|
42
42
|
"pytest>=8.3.4",
|
|
43
43
|
"pytest-asyncio>=0.26.0",
|
|
44
44
|
"pytest-cov>=6.1.0",
|
|
45
|
-
"typos>=1.31.
|
|
46
|
-
"ruff>=0.11.
|
|
45
|
+
"typos>=1.31.2",
|
|
46
|
+
"ruff>=0.11.7",
|
|
47
47
|
"mypy>=1.15.0",
|
|
48
|
-
"pyright>=1.1.
|
|
48
|
+
"pyright>=1.1.400",
|
|
49
49
|
"pre-commit>=4.2.0",
|
|
50
50
|
"bump2version>=1.0.1",
|
|
51
51
|
"build>=1.2.2",
|
|
@@ -244,7 +244,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer):
|
|
|
244
244
|
# Get input schema if available
|
|
245
245
|
try:
|
|
246
246
|
# Convert schema to proto format
|
|
247
|
-
input_schema_proto = self.module_class.get_input_format(request.llm_format)
|
|
247
|
+
input_schema_proto = self.module_class.get_input_format(llm_format=request.llm_format)
|
|
248
248
|
input_format_struct = json_format.Parse(
|
|
249
249
|
text=input_schema_proto,
|
|
250
250
|
message=struct_pb2.Struct(), # pylint: disable=no-member
|
|
@@ -280,7 +280,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer):
|
|
|
280
280
|
# Get output schema if available
|
|
281
281
|
try:
|
|
282
282
|
# Convert schema to proto format
|
|
283
|
-
output_schema_proto = self.module_class.get_output_format(request.llm_format)
|
|
283
|
+
output_schema_proto = self.module_class.get_output_format(llm_format=request.llm_format)
|
|
284
284
|
output_format_struct = json_format.Parse(
|
|
285
285
|
text=output_schema_proto,
|
|
286
286
|
message=struct_pb2.Struct(), # pylint: disable=no-member
|
|
@@ -316,7 +316,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer):
|
|
|
316
316
|
# Get setup schema if available
|
|
317
317
|
try:
|
|
318
318
|
# Convert schema to proto format
|
|
319
|
-
setup_schema_proto = self.module_class.get_setup_format(request.llm_format)
|
|
319
|
+
setup_schema_proto = self.module_class.get_setup_format(llm_format=request.llm_format)
|
|
320
320
|
setup_format_struct = json_format.Parse(
|
|
321
321
|
text=setup_schema_proto,
|
|
322
322
|
message=struct_pb2.Struct(), # pylint: disable=no-member
|
|
@@ -332,3 +332,39 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer):
|
|
|
332
332
|
success=True,
|
|
333
333
|
setup_schema=setup_format_struct,
|
|
334
334
|
)
|
|
335
|
+
|
|
336
|
+
def GetModuleSecret( # noqa: N802
|
|
337
|
+
self,
|
|
338
|
+
request: information_pb2.GetModuleSecretRequest,
|
|
339
|
+
context: grpc.ServicerContext,
|
|
340
|
+
) -> information_pb2.GetModuleSecretResponse:
|
|
341
|
+
"""Get information about the module's secrets.
|
|
342
|
+
|
|
343
|
+
Args:
|
|
344
|
+
request: The get module secret request.
|
|
345
|
+
context: The gRPC context.
|
|
346
|
+
|
|
347
|
+
Returns:
|
|
348
|
+
A response with the module's secret schema.
|
|
349
|
+
"""
|
|
350
|
+
logger.info("GetModuleSecret called for module: '%s'", self.module_class.__name__)
|
|
351
|
+
|
|
352
|
+
# Get secret schema if available
|
|
353
|
+
try:
|
|
354
|
+
# Convert schema to proto format
|
|
355
|
+
secret_schema_proto = self.module_class.get_secret_format(llm_format=request.llm_format)
|
|
356
|
+
secret_format_struct = json_format.Parse(
|
|
357
|
+
text=secret_schema_proto,
|
|
358
|
+
message=struct_pb2.Struct(), # pylint: disable=no-member
|
|
359
|
+
ignore_unknown_fields=True,
|
|
360
|
+
)
|
|
361
|
+
except NotImplementedError as e:
|
|
362
|
+
logger.warning(e)
|
|
363
|
+
context.set_code(grpc.StatusCode.UNIMPLEMENTED)
|
|
364
|
+
context.set_details(e)
|
|
365
|
+
return information_pb2.GetModuleSecretResponse()
|
|
366
|
+
|
|
367
|
+
return information_pb2.GetModuleSecretResponse(
|
|
368
|
+
success=True,
|
|
369
|
+
secret_schema=secret_format_struct,
|
|
370
|
+
)
|
|
@@ -17,6 +17,7 @@ from digitalkin.services.registry.registry_strategy import RegistryStrategy
|
|
|
17
17
|
from digitalkin.services.services_config import ServicesConfig, ServicesStrategy
|
|
18
18
|
from digitalkin.services.snapshot.snapshot_strategy import SnapshotStrategy
|
|
19
19
|
from digitalkin.services.storage.storage_strategy import StorageStrategy
|
|
20
|
+
from digitalkin.utils.llm_ready_schema import llm_ready_schema
|
|
20
21
|
|
|
21
22
|
|
|
22
23
|
class BaseModule(ABC, Generic[InputModelT, OutputModelT, SetupModelT, SecretModelT]):
|
|
@@ -73,7 +74,7 @@ class BaseModule(ABC, Generic[InputModelT, OutputModelT, SetupModelT, SecretMode
|
|
|
73
74
|
return self._status
|
|
74
75
|
|
|
75
76
|
@classmethod
|
|
76
|
-
def get_secret_format(cls, llm_format: bool) -> str:
|
|
77
|
+
def get_secret_format(cls, *, llm_format: bool) -> str:
|
|
77
78
|
"""Get the JSON schema of the secret format model.
|
|
78
79
|
|
|
79
80
|
Raises:
|
|
@@ -84,13 +85,13 @@ class BaseModule(ABC, Generic[InputModelT, OutputModelT, SetupModelT, SecretMode
|
|
|
84
85
|
"""
|
|
85
86
|
if cls.secret_format is not None:
|
|
86
87
|
if llm_format:
|
|
87
|
-
return json.dumps(cls.secret_format, indent=2)
|
|
88
|
+
return json.dumps(llm_ready_schema(cls.secret_format), indent=2)
|
|
88
89
|
return json.dumps(cls.secret_format.model_json_schema(), indent=2)
|
|
89
90
|
msg = f"{cls.__name__}' class does not define a 'secret_format'."
|
|
90
91
|
raise NotImplementedError(msg)
|
|
91
92
|
|
|
92
93
|
@classmethod
|
|
93
|
-
def get_input_format(cls, llm_format: bool) -> str:
|
|
94
|
+
def get_input_format(cls, *, llm_format: bool) -> str:
|
|
94
95
|
"""Get the JSON schema of the input format model.
|
|
95
96
|
|
|
96
97
|
Raises:
|
|
@@ -101,13 +102,13 @@ class BaseModule(ABC, Generic[InputModelT, OutputModelT, SetupModelT, SecretMode
|
|
|
101
102
|
"""
|
|
102
103
|
if cls.input_format is not None:
|
|
103
104
|
if llm_format:
|
|
104
|
-
return json.dumps(cls.input_format, indent=2)
|
|
105
|
+
return json.dumps(llm_ready_schema(cls.input_format), indent=2)
|
|
105
106
|
return json.dumps(cls.input_format.model_json_schema(), indent=2)
|
|
106
107
|
msg = f"{cls.__name__}' class does not define an 'input_format'."
|
|
107
108
|
raise NotImplementedError(msg)
|
|
108
109
|
|
|
109
110
|
@classmethod
|
|
110
|
-
def get_output_format(cls, llm_format: bool) -> str:
|
|
111
|
+
def get_output_format(cls, *, llm_format: bool) -> str:
|
|
111
112
|
"""Get the JSON schema of the output format model.
|
|
112
113
|
|
|
113
114
|
Raises:
|
|
@@ -118,13 +119,13 @@ class BaseModule(ABC, Generic[InputModelT, OutputModelT, SetupModelT, SecretMode
|
|
|
118
119
|
"""
|
|
119
120
|
if cls.output_format is not None:
|
|
120
121
|
if llm_format:
|
|
121
|
-
return json.dumps(cls.output_format, indent=2)
|
|
122
|
+
return json.dumps(llm_ready_schema(cls.output_format), indent=2)
|
|
122
123
|
return json.dumps(cls.output_format.model_json_schema(), indent=2)
|
|
123
124
|
msg = "'%s' class does not define an 'output_format'."
|
|
124
125
|
raise NotImplementedError(msg)
|
|
125
126
|
|
|
126
127
|
@classmethod
|
|
127
|
-
def get_setup_format(cls, llm_format: bool) -> str:
|
|
128
|
+
def get_setup_format(cls, *, llm_format: bool) -> str:
|
|
128
129
|
"""Gets the JSON schema of the setup format model.
|
|
129
130
|
|
|
130
131
|
Raises:
|
|
@@ -135,7 +136,7 @@ class BaseModule(ABC, Generic[InputModelT, OutputModelT, SetupModelT, SecretMode
|
|
|
135
136
|
"""
|
|
136
137
|
if cls.setup_format is not None:
|
|
137
138
|
if llm_format:
|
|
138
|
-
return json.dumps(cls.setup_format, indent=2)
|
|
139
|
+
return json.dumps(llm_ready_schema(cls.setup_format), indent=2)
|
|
139
140
|
return json.dumps(cls.setup_format.model_json_schema(), indent=2)
|
|
140
141
|
msg = "'%s' class does not define an 'setup_format'."
|
|
141
142
|
raise NotImplementedError(msg)
|
|
@@ -0,0 +1,220 @@
|
|
|
1
|
+
"""This module implements the default storage strategy."""
|
|
2
|
+
|
|
3
|
+
import datetime
|
|
4
|
+
import json
|
|
5
|
+
import logging
|
|
6
|
+
import tempfile
|
|
7
|
+
from pathlib import Path
|
|
8
|
+
from typing import Any
|
|
9
|
+
|
|
10
|
+
from pydantic import BaseModel
|
|
11
|
+
|
|
12
|
+
from digitalkin.services.storage.storage_strategy import DataType, StorageRecord, StorageStrategy
|
|
13
|
+
|
|
14
|
+
logger = logging.getLogger(__name__)
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
class DefaultStorage(StorageStrategy):
|
|
18
|
+
"""Persist records in a local JSON file for quick local development.
|
|
19
|
+
|
|
20
|
+
File format: a JSON object of
|
|
21
|
+
{ "<collection>:<record_id>": { ... StorageRecord fields ... },
|
|
22
|
+
"""
|
|
23
|
+
|
|
24
|
+
@staticmethod
|
|
25
|
+
def _json_default(o: Any) -> str: # noqa: ANN401
|
|
26
|
+
"""JSON serializer for non-standard types (datetime → ISO).
|
|
27
|
+
|
|
28
|
+
Args:
|
|
29
|
+
o: The object to serialize
|
|
30
|
+
|
|
31
|
+
Returns:
|
|
32
|
+
str: The serialized object
|
|
33
|
+
|
|
34
|
+
Raises:
|
|
35
|
+
TypeError: If the object is not serializable
|
|
36
|
+
"""
|
|
37
|
+
if isinstance(o, datetime.datetime):
|
|
38
|
+
return o.isoformat()
|
|
39
|
+
msg = f"Type {o.__class__.__name__} not serializable"
|
|
40
|
+
raise TypeError(msg)
|
|
41
|
+
|
|
42
|
+
def _load_from_file(self) -> dict[str, StorageRecord]:
|
|
43
|
+
"""Load storage data from the file.
|
|
44
|
+
|
|
45
|
+
Returns:
|
|
46
|
+
A dictionary containing the loaded storage records
|
|
47
|
+
"""
|
|
48
|
+
if not self.storage_file.exists():
|
|
49
|
+
return {}
|
|
50
|
+
|
|
51
|
+
try:
|
|
52
|
+
raw = json.loads(self.storage_file.read_text(encoding="utf-8"))
|
|
53
|
+
out: dict[str, StorageRecord] = {}
|
|
54
|
+
|
|
55
|
+
for key, rd in raw.items():
|
|
56
|
+
# rd is a dict with the StorageRecord fields
|
|
57
|
+
model_cls = self.config.get(rd["collection"])
|
|
58
|
+
if not model_cls:
|
|
59
|
+
logger.warning("No model for collection '%s'", rd["collection"])
|
|
60
|
+
continue
|
|
61
|
+
data_model = model_cls.model_validate(rd["data"])
|
|
62
|
+
rec = StorageRecord(
|
|
63
|
+
mission_id=rd["mission_id"],
|
|
64
|
+
collection=rd["collection"],
|
|
65
|
+
record_id=rd["record_id"],
|
|
66
|
+
data=data_model,
|
|
67
|
+
data_type=DataType[rd["data_type"]],
|
|
68
|
+
creation_date=datetime.datetime.fromisoformat(rd["creation_date"])
|
|
69
|
+
if rd.get("creation_date")
|
|
70
|
+
else None,
|
|
71
|
+
update_date=datetime.datetime.fromisoformat(rd["update_date"]) if rd.get("update_date") else None,
|
|
72
|
+
)
|
|
73
|
+
out[key] = rec
|
|
74
|
+
except Exception:
|
|
75
|
+
logger.exception("Failed to load default storage file")
|
|
76
|
+
return {}
|
|
77
|
+
return out
|
|
78
|
+
|
|
79
|
+
def _save_to_file(self) -> None:
|
|
80
|
+
"""Atomically write `self.storage` back to disk as JSON."""
|
|
81
|
+
self.storage_file.parent.mkdir(parents=True, exist_ok=True)
|
|
82
|
+
with tempfile.NamedTemporaryFile(
|
|
83
|
+
mode="w", encoding="utf-8", delete=False, dir=str(self.storage_file.parent), suffix=".tmp"
|
|
84
|
+
) as temp:
|
|
85
|
+
try:
|
|
86
|
+
# Convert storage to a serializable format
|
|
87
|
+
serial: dict[str, dict] = {}
|
|
88
|
+
for key, record in self.storage.items():
|
|
89
|
+
serial[key] = {
|
|
90
|
+
"mission_id": record.mission_id,
|
|
91
|
+
"collection": record.collection,
|
|
92
|
+
"record_id": record.record_id,
|
|
93
|
+
"data_type": record.data_type.name,
|
|
94
|
+
"data": record.data.model_dump(),
|
|
95
|
+
"creation_date": record.creation_date.isoformat() if record.creation_date else None,
|
|
96
|
+
"update_date": record.update_date.isoformat() if record.update_date else None,
|
|
97
|
+
}
|
|
98
|
+
json.dump(serial, temp, indent=2, default=self._json_default)
|
|
99
|
+
temp.flush()
|
|
100
|
+
Path(temp.name).replace(self.storage_file)
|
|
101
|
+
except Exception:
|
|
102
|
+
logger.exception("Unexpected error saving storage")
|
|
103
|
+
|
|
104
|
+
def _store(self, record: StorageRecord) -> StorageRecord:
|
|
105
|
+
"""Store a new record in the database and persist to file.
|
|
106
|
+
|
|
107
|
+
Args:
|
|
108
|
+
record: The record to store
|
|
109
|
+
|
|
110
|
+
Returns:
|
|
111
|
+
str: The ID of the new record
|
|
112
|
+
|
|
113
|
+
Raises:
|
|
114
|
+
ValueError: If the record already exists
|
|
115
|
+
"""
|
|
116
|
+
key = f"{record.collection}:{record.record_id}"
|
|
117
|
+
if key in self.storage:
|
|
118
|
+
msg = f"Document {key!r} already exists"
|
|
119
|
+
raise ValueError(msg)
|
|
120
|
+
now = datetime.datetime.now(datetime.timezone.utc)
|
|
121
|
+
record.creation_date = now
|
|
122
|
+
record.update_date = now
|
|
123
|
+
self.storage[key] = record
|
|
124
|
+
self._save_to_file()
|
|
125
|
+
logger.debug("Created %s", key)
|
|
126
|
+
return record
|
|
127
|
+
|
|
128
|
+
def _read(self, collection: str, record_id: str) -> StorageRecord | None:
|
|
129
|
+
"""Get records from the database.
|
|
130
|
+
|
|
131
|
+
Args:
|
|
132
|
+
collection: The unique name to retrieve data for
|
|
133
|
+
record_id: The unique ID of the record
|
|
134
|
+
|
|
135
|
+
Returns:
|
|
136
|
+
StorageRecord: The corresponding record
|
|
137
|
+
"""
|
|
138
|
+
key = f"{collection}:{record_id}"
|
|
139
|
+
return self.storage.get(key)
|
|
140
|
+
|
|
141
|
+
def _update(self, collection: str, record_id: str, data: BaseModel) -> StorageRecord | None:
|
|
142
|
+
"""Update records in the database and persist to file.
|
|
143
|
+
|
|
144
|
+
Args:
|
|
145
|
+
collection: The unique name to retrieve data for
|
|
146
|
+
record_id: The unique ID of the record
|
|
147
|
+
data: The data to modify
|
|
148
|
+
|
|
149
|
+
Returns:
|
|
150
|
+
StorageRecord: The modified record
|
|
151
|
+
"""
|
|
152
|
+
key = f"{collection}:{record_id}"
|
|
153
|
+
rec = self.storage.get(key)
|
|
154
|
+
if not rec:
|
|
155
|
+
return None
|
|
156
|
+
rec.data = data
|
|
157
|
+
rec.update_date = datetime.datetime.now(datetime.timezone.utc)
|
|
158
|
+
self._save_to_file()
|
|
159
|
+
logger.debug("Modified %s", key)
|
|
160
|
+
return rec
|
|
161
|
+
|
|
162
|
+
def _remove(self, collection: str, record_id: str) -> bool:
|
|
163
|
+
"""Delete records from the database and update file.
|
|
164
|
+
|
|
165
|
+
Args:
|
|
166
|
+
collection: The unique name to retrieve data for
|
|
167
|
+
record_id: The unique ID of the record
|
|
168
|
+
|
|
169
|
+
Returns:
|
|
170
|
+
bool: True if the record was removed, False otherwise
|
|
171
|
+
"""
|
|
172
|
+
key = f"{collection}:{record_id}"
|
|
173
|
+
if key not in self.storage:
|
|
174
|
+
return False
|
|
175
|
+
del self.storage[key]
|
|
176
|
+
self._save_to_file()
|
|
177
|
+
logger.debug("Removed %s", key)
|
|
178
|
+
return True
|
|
179
|
+
|
|
180
|
+
def _list(self, collection: str) -> list[StorageRecord]:
|
|
181
|
+
"""Implements StorageStrategy._list.
|
|
182
|
+
|
|
183
|
+
Args:
|
|
184
|
+
collection: The unique name to retrieve data for
|
|
185
|
+
|
|
186
|
+
Returns:
|
|
187
|
+
A list of storage records
|
|
188
|
+
"""
|
|
189
|
+
prefix = f"{collection}:"
|
|
190
|
+
return [r for k, r in self.storage.items() if k.startswith(prefix)]
|
|
191
|
+
|
|
192
|
+
def _remove_collection(self, collection: str) -> bool:
|
|
193
|
+
"""Implements StorageStrategy._remove_collection.
|
|
194
|
+
|
|
195
|
+
Args:
|
|
196
|
+
collection: The unique name to retrieve data for
|
|
197
|
+
|
|
198
|
+
Returns:
|
|
199
|
+
bool: True if the collection was removed, False otherwise
|
|
200
|
+
"""
|
|
201
|
+
prefix = f"{collection}:"
|
|
202
|
+
to_delete = [k for k in self.storage if k.startswith(prefix)]
|
|
203
|
+
for k in to_delete:
|
|
204
|
+
del self.storage[k]
|
|
205
|
+
self._save_to_file()
|
|
206
|
+
logger.debug("Removed collection %s (%d docs)", collection, len(to_delete))
|
|
207
|
+
return True
|
|
208
|
+
|
|
209
|
+
def __init__(
|
|
210
|
+
self,
|
|
211
|
+
mission_id: str,
|
|
212
|
+
config: dict[str, type[BaseModel]],
|
|
213
|
+
storage_file_path: str = "local_storage",
|
|
214
|
+
**kwargs, # noqa: ANN003, ARG002
|
|
215
|
+
) -> None:
|
|
216
|
+
"""Initialize the storage."""
|
|
217
|
+
super().__init__(mission_id=mission_id, config=config)
|
|
218
|
+
self.storage_file_path = f"{self.mission_id}_{storage_file_path}.json"
|
|
219
|
+
self.storage_file = Path(self.storage_file_path)
|
|
220
|
+
self.storage = self._load_from_file()
|
|
@@ -0,0 +1,197 @@
|
|
|
1
|
+
"""This module implements the default storage strategy."""
|
|
2
|
+
|
|
3
|
+
import logging
|
|
4
|
+
|
|
5
|
+
from digitalkin_proto.digitalkin.storage.v2 import data_pb2, storage_service_pb2_grpc
|
|
6
|
+
from google.protobuf import json_format
|
|
7
|
+
from google.protobuf.struct_pb2 import Struct
|
|
8
|
+
from pydantic import BaseModel
|
|
9
|
+
|
|
10
|
+
from digitalkin.grpc_servers.utils.grpc_client_wrapper import GrpcClientWrapper
|
|
11
|
+
from digitalkin.grpc_servers.utils.models import ClientConfig
|
|
12
|
+
from digitalkin.services.storage.storage_strategy import DataType, StorageRecord, StorageServiceError, StorageStrategy
|
|
13
|
+
|
|
14
|
+
logger = logging.getLogger(__name__)
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
class GrpcStorage(StorageStrategy, GrpcClientWrapper):
|
|
18
|
+
"""This class implements the default storage strategy."""
|
|
19
|
+
|
|
20
|
+
def _build_record_from_proto(self, proto: data_pb2.StorageRecord) -> StorageRecord:
|
|
21
|
+
"""Convert a protobuf StorageRecord message into our Pydantic model.
|
|
22
|
+
|
|
23
|
+
Args:
|
|
24
|
+
proto: gRPC StorageRecord
|
|
25
|
+
|
|
26
|
+
Returns:
|
|
27
|
+
A fully validated StorageRecord.
|
|
28
|
+
"""
|
|
29
|
+
raw = json_format.MessageToDict(
|
|
30
|
+
proto,
|
|
31
|
+
preserving_proto_field_name=True,
|
|
32
|
+
always_print_fields_with_no_presence=True,
|
|
33
|
+
)
|
|
34
|
+
mission = raw["mission_id"]
|
|
35
|
+
coll = raw["collection"]
|
|
36
|
+
rid = raw["record_id"]
|
|
37
|
+
dtype = DataType[raw["data_type"]]
|
|
38
|
+
payload = raw.get("data", {})
|
|
39
|
+
|
|
40
|
+
validated = self._validate_data(coll, payload)
|
|
41
|
+
return StorageRecord(
|
|
42
|
+
mission_id=mission,
|
|
43
|
+
collection=coll,
|
|
44
|
+
record_id=rid,
|
|
45
|
+
data=validated,
|
|
46
|
+
data_type=dtype,
|
|
47
|
+
creation_date=raw.get("creation_date"),
|
|
48
|
+
update_date=raw.get("update_date"),
|
|
49
|
+
)
|
|
50
|
+
|
|
51
|
+
def _store(self, record: StorageRecord) -> StorageRecord:
|
|
52
|
+
"""Create a new record in the database.
|
|
53
|
+
|
|
54
|
+
Parameters:
|
|
55
|
+
record: The record to store
|
|
56
|
+
|
|
57
|
+
Returns:
|
|
58
|
+
StorageRecord: The corresponding record
|
|
59
|
+
|
|
60
|
+
Raises:
|
|
61
|
+
StorageServiceError: If there is an error while storing the record
|
|
62
|
+
"""
|
|
63
|
+
try:
|
|
64
|
+
data_struct = Struct()
|
|
65
|
+
data_struct.update(record.data.model_dump())
|
|
66
|
+
req = data_pb2.StoreRecordRequest(
|
|
67
|
+
data=data_struct,
|
|
68
|
+
mission_id=record.mission_id,
|
|
69
|
+
collection=record.collection,
|
|
70
|
+
record_id=record.record_id,
|
|
71
|
+
data_type=record.data_type.name,
|
|
72
|
+
)
|
|
73
|
+
resp = self.exec_grpc_query("StoreRecord", req)
|
|
74
|
+
return self._build_record_from_proto(resp.stored_data)
|
|
75
|
+
except Exception as e:
|
|
76
|
+
logger.exception("gRPC StoreRecord failed for %s:%s", record.collection, record.record_id)
|
|
77
|
+
raise StorageServiceError(str(e)) from e
|
|
78
|
+
|
|
79
|
+
def _read(self, collection: str, record_id: str) -> StorageRecord | None:
|
|
80
|
+
"""Fetch a single document by collection + record_id.
|
|
81
|
+
|
|
82
|
+
Returns:
|
|
83
|
+
StorageData: The record
|
|
84
|
+
"""
|
|
85
|
+
try:
|
|
86
|
+
req = data_pb2.ReadRecordRequest(
|
|
87
|
+
mission_id=self.mission_id,
|
|
88
|
+
collection=collection,
|
|
89
|
+
record_id=record_id,
|
|
90
|
+
)
|
|
91
|
+
resp = self.exec_grpc_query("ReadRecord", req)
|
|
92
|
+
return self._build_record_from_proto(resp.stored_data)
|
|
93
|
+
except Exception:
|
|
94
|
+
logger.exception("gRPC ReadRecord failed for %s:%s", collection, record_id)
|
|
95
|
+
return None
|
|
96
|
+
|
|
97
|
+
def _update(self, collection: str, record_id: str, data: BaseModel) -> StorageRecord | None:
|
|
98
|
+
"""Overwrite a document via gRPC.
|
|
99
|
+
|
|
100
|
+
Args:
|
|
101
|
+
collection: The unique name for the record type
|
|
102
|
+
record_id: The unique ID for the record
|
|
103
|
+
data: The validated data model
|
|
104
|
+
|
|
105
|
+
Returns:
|
|
106
|
+
StorageRecord: The updated record
|
|
107
|
+
"""
|
|
108
|
+
try:
|
|
109
|
+
struct = Struct()
|
|
110
|
+
struct.update(data.model_dump())
|
|
111
|
+
req = data_pb2.UpdateRecordRequest(
|
|
112
|
+
data=struct,
|
|
113
|
+
mission_id=self.mission_id,
|
|
114
|
+
collection=collection,
|
|
115
|
+
record_id=record_id,
|
|
116
|
+
)
|
|
117
|
+
resp = self.exec_grpc_query("ModifyRecord", req)
|
|
118
|
+
return self._build_record_from_proto(resp.stored_data)
|
|
119
|
+
except Exception:
|
|
120
|
+
logger.exception("gRPC ModifyRecord failed for %s:%s", collection, record_id)
|
|
121
|
+
return None
|
|
122
|
+
|
|
123
|
+
def _remove(self, collection: str, record_id: str) -> bool:
|
|
124
|
+
"""Delete a document via gRPC.
|
|
125
|
+
|
|
126
|
+
Args:
|
|
127
|
+
collection: The unique name for the record type
|
|
128
|
+
record_id: The unique ID for the record
|
|
129
|
+
|
|
130
|
+
Returns:
|
|
131
|
+
bool: True if the record was deleted, False otherwise
|
|
132
|
+
"""
|
|
133
|
+
try:
|
|
134
|
+
req = data_pb2.RemoveRecordRequest(
|
|
135
|
+
mission_id=self.mission_id,
|
|
136
|
+
collection=collection,
|
|
137
|
+
record_id=record_id,
|
|
138
|
+
)
|
|
139
|
+
self.exec_grpc_query("RemoveRecord", req)
|
|
140
|
+
except Exception:
|
|
141
|
+
logger.exception("gRPC RemoveRecord failed for %s:%s", collection, record_id)
|
|
142
|
+
return False
|
|
143
|
+
return True
|
|
144
|
+
|
|
145
|
+
def _list(self, collection: str) -> list[StorageRecord]:
|
|
146
|
+
"""List all documents in a collection via gRPC.
|
|
147
|
+
|
|
148
|
+
Args:
|
|
149
|
+
collection: The unique name for the record type
|
|
150
|
+
|
|
151
|
+
Returns:
|
|
152
|
+
list[StorageRecord]: A list of storage records
|
|
153
|
+
"""
|
|
154
|
+
try:
|
|
155
|
+
req = data_pb2.ListRecordsRequest(
|
|
156
|
+
mission_id=self.mission_id,
|
|
157
|
+
collection=collection,
|
|
158
|
+
)
|
|
159
|
+
resp = self.exec_grpc_query("ListRecords", req)
|
|
160
|
+
return [self._build_record_from_proto(r) for r in resp.records]
|
|
161
|
+
except Exception:
|
|
162
|
+
logger.exception("gRPC ListRecords failed for %s", collection)
|
|
163
|
+
return []
|
|
164
|
+
|
|
165
|
+
def _remove_collection(self, collection: str) -> bool:
|
|
166
|
+
"""Delete an entire collection via gRPC.
|
|
167
|
+
|
|
168
|
+
Args:
|
|
169
|
+
collection: The unique name for the record type
|
|
170
|
+
|
|
171
|
+
Returns:
|
|
172
|
+
bool: True if the collection was deleted, False otherwise
|
|
173
|
+
"""
|
|
174
|
+
try:
|
|
175
|
+
req = data_pb2.RemoveCollectionRequest(
|
|
176
|
+
mission_id=self.mission_id,
|
|
177
|
+
collection=collection,
|
|
178
|
+
)
|
|
179
|
+
self.exec_grpc_query("RemoveCollection", req)
|
|
180
|
+
except Exception:
|
|
181
|
+
logger.exception("gRPC RemoveCollection failed for %s", collection)
|
|
182
|
+
return False
|
|
183
|
+
return True
|
|
184
|
+
|
|
185
|
+
def __init__(
|
|
186
|
+
self,
|
|
187
|
+
mission_id: str,
|
|
188
|
+
config: dict[str, type[BaseModel]],
|
|
189
|
+
client_config: ClientConfig,
|
|
190
|
+
**kwargs, # noqa: ANN003, ARG002
|
|
191
|
+
) -> None:
|
|
192
|
+
"""Initialize the storage."""
|
|
193
|
+
super().__init__(mission_id=mission_id, config=config)
|
|
194
|
+
|
|
195
|
+
channel = self._init_channel(client_config)
|
|
196
|
+
self.stub = storage_service_pb2_grpc.StorageServiceStub(channel)
|
|
197
|
+
logger.debug("Channel client 'storage' initialized succesfully")
|