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.
Files changed (86) hide show
  1. {digitalkin-0.2.6 → digitalkin-0.2.8}/PKG-INFO +7 -7
  2. {digitalkin-0.2.6 → digitalkin-0.2.8}/examples/modules/storage_module.py +9 -4
  3. {digitalkin-0.2.6 → digitalkin-0.2.8}/pyproject.toml +7 -7
  4. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/__version__.py +1 -1
  5. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/grpc_servers/module_servicer.py +39 -3
  6. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/modules/_base_module.py +9 -8
  7. digitalkin-0.2.8/src/digitalkin/services/storage/default_storage.py +220 -0
  8. digitalkin-0.2.8/src/digitalkin/services/storage/grpc_storage.py +197 -0
  9. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/storage/storage_strategy.py +156 -104
  10. digitalkin-0.2.8/src/digitalkin/utils/llm_ready_schema.py +75 -0
  11. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin.egg-info/PKG-INFO +7 -7
  12. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin.egg-info/SOURCES.txt +2 -1
  13. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin.egg-info/requires.txt +6 -6
  14. digitalkin-0.2.6/src/digitalkin/services/storage/default_storage.py +0 -218
  15. digitalkin-0.2.6/src/digitalkin/services/storage/grpc_storage.py +0 -156
  16. {digitalkin-0.2.6 → digitalkin-0.2.8}/LICENSE +0 -0
  17. {digitalkin-0.2.6 → digitalkin-0.2.8}/README.md +0 -0
  18. {digitalkin-0.2.6 → digitalkin-0.2.8}/examples/base_server/__init__.py +0 -0
  19. {digitalkin-0.2.6 → digitalkin-0.2.8}/examples/base_server/mock/__init__.py +0 -0
  20. {digitalkin-0.2.6 → digitalkin-0.2.8}/examples/base_server/mock/mock_pb2.py +0 -0
  21. {digitalkin-0.2.6 → digitalkin-0.2.8}/examples/base_server/mock/mock_pb2_grpc.py +0 -0
  22. {digitalkin-0.2.6 → digitalkin-0.2.8}/examples/base_server/server_async_insecure.py +0 -0
  23. {digitalkin-0.2.6 → digitalkin-0.2.8}/examples/base_server/server_async_secure.py +0 -0
  24. {digitalkin-0.2.6 → digitalkin-0.2.8}/examples/base_server/server_sync_insecure.py +0 -0
  25. {digitalkin-0.2.6 → digitalkin-0.2.8}/examples/base_server/server_sync_secure.py +0 -0
  26. {digitalkin-0.2.6 → digitalkin-0.2.8}/examples/modules/__init__.py +0 -0
  27. {digitalkin-0.2.6 → digitalkin-0.2.8}/examples/modules/minimal_llm_module.py +0 -0
  28. {digitalkin-0.2.6 → digitalkin-0.2.8}/examples/modules/text_transform_module.py +0 -0
  29. {digitalkin-0.2.6 → digitalkin-0.2.8}/setup.cfg +0 -0
  30. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/__init__.py +0 -0
  31. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/grpc_servers/__init__.py +0 -0
  32. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/grpc_servers/_base_server.py +0 -0
  33. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/grpc_servers/module_server.py +0 -0
  34. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/grpc_servers/registry_server.py +0 -0
  35. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/grpc_servers/registry_servicer.py +0 -0
  36. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/grpc_servers/utils/exceptions.py +0 -0
  37. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/grpc_servers/utils/factory.py +0 -0
  38. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/grpc_servers/utils/grpc_client_wrapper.py +0 -0
  39. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/grpc_servers/utils/models.py +0 -0
  40. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/grpc_servers/utils/types.py +0 -0
  41. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/logger.py +0 -0
  42. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/models/__init__.py +0 -0
  43. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/models/module/__init__.py +0 -0
  44. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/models/module/module.py +0 -0
  45. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/models/module/module_types.py +0 -0
  46. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/models/services/__init__.py +0 -0
  47. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/models/services/cost.py +0 -0
  48. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/models/services/storage.py +0 -0
  49. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/modules/__init__.py +0 -0
  50. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/modules/archetype_module.py +0 -0
  51. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/modules/job_manager.py +0 -0
  52. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/modules/tool_module.py +0 -0
  53. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/modules/trigger_module.py +0 -0
  54. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/py.typed +0 -0
  55. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/__init__.py +0 -0
  56. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/agent/__init__.py +0 -0
  57. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/agent/agent_strategy.py +0 -0
  58. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/agent/default_agent.py +0 -0
  59. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/base_strategy.py +0 -0
  60. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/cost/__init__.py +0 -0
  61. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/cost/cost_strategy.py +0 -0
  62. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/cost/default_cost.py +0 -0
  63. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/cost/grpc_cost.py +0 -0
  64. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/filesystem/__init__.py +0 -0
  65. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/filesystem/default_filesystem.py +0 -0
  66. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/filesystem/filesystem_strategy.py +0 -0
  67. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/filesystem/grpc_filesystem.py +0 -0
  68. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/identity/__init__.py +0 -0
  69. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/identity/default_identity.py +0 -0
  70. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/identity/identity_strategy.py +0 -0
  71. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/registry/__init__.py +0 -0
  72. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/registry/default_registry.py +0 -0
  73. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/registry/registry_strategy.py +0 -0
  74. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/services_config.py +0 -0
  75. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/services_models.py +0 -0
  76. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/setup/default_setup.py +0 -0
  77. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/setup/grpc_setup.py +0 -0
  78. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/setup/setup_strategy.py +0 -0
  79. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/snapshot/__init__.py +0 -0
  80. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/snapshot/default_snapshot.py +0 -0
  81. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/snapshot/snapshot_strategy.py +0 -0
  82. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/services/storage/__init__.py +0 -0
  83. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/utils/__init__.py +0 -0
  84. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin/utils/arg_parser.py +0 -0
  85. {digitalkin-0.2.6 → digitalkin-0.2.8}/src/digitalkin.egg-info/dependency_links.txt +0 -0
  86. {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.6
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.7
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.70.0
460
- Requires-Dist: pydantic>=2.11.3
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.1; extra == "dev"
466
- Requires-Dist: ruff>=0.11.2; extra == "dev"
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.398; extra == "dev"
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("example_outputs", output_data.model_dump(), data_type="OUTPUT")
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.6"
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.7",
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.70.0",
37
- "pydantic>=2.11.3",
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.1",
46
- "ruff>=0.11.2",
45
+ "typos>=1.31.2",
46
+ "ruff>=0.11.7",
47
47
  "mypy>=1.15.0",
48
- "pyright>=1.1.398",
48
+ "pyright>=1.1.400",
49
49
  "pre-commit>=4.2.0",
50
50
  "bump2version>=1.0.1",
51
51
  "build>=1.2.2",
@@ -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.6"
8
+ __version__ = "0.2.8"
@@ -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: # noqa: FBT001
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: # noqa: FBT001
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: # noqa: FBT001
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: # noqa: FBT001
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")