digitalkin 0.2.5__tar.gz → 0.2.7__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 (84) hide show
  1. {digitalkin-0.2.5 → digitalkin-0.2.7}/PKG-INFO +1 -1
  2. {digitalkin-0.2.5 → digitalkin-0.2.7}/examples/modules/minimal_llm_module.py +4 -5
  3. {digitalkin-0.2.5 → digitalkin-0.2.7}/examples/modules/text_transform_module.py +8 -5
  4. {digitalkin-0.2.5 → digitalkin-0.2.7}/pyproject.toml +1 -1
  5. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/__version__.py +1 -1
  6. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/grpc_servers/module_servicer.py +39 -3
  7. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/modules/_base_module.py +9 -8
  8. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/services/filesystem/grpc_filesystem.py +4 -4
  9. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/services/storage/grpc_storage.py +38 -5
  10. digitalkin-0.2.7/src/digitalkin/utils/llm_ready_schema.py +75 -0
  11. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin.egg-info/PKG-INFO +1 -1
  12. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin.egg-info/SOURCES.txt +2 -1
  13. {digitalkin-0.2.5 → digitalkin-0.2.7}/LICENSE +0 -0
  14. {digitalkin-0.2.5 → digitalkin-0.2.7}/README.md +0 -0
  15. {digitalkin-0.2.5 → digitalkin-0.2.7}/examples/base_server/__init__.py +0 -0
  16. {digitalkin-0.2.5 → digitalkin-0.2.7}/examples/base_server/mock/__init__.py +0 -0
  17. {digitalkin-0.2.5 → digitalkin-0.2.7}/examples/base_server/mock/mock_pb2.py +0 -0
  18. {digitalkin-0.2.5 → digitalkin-0.2.7}/examples/base_server/mock/mock_pb2_grpc.py +0 -0
  19. {digitalkin-0.2.5 → digitalkin-0.2.7}/examples/base_server/server_async_insecure.py +0 -0
  20. {digitalkin-0.2.5 → digitalkin-0.2.7}/examples/base_server/server_async_secure.py +0 -0
  21. {digitalkin-0.2.5 → digitalkin-0.2.7}/examples/base_server/server_sync_insecure.py +0 -0
  22. {digitalkin-0.2.5 → digitalkin-0.2.7}/examples/base_server/server_sync_secure.py +0 -0
  23. {digitalkin-0.2.5 → digitalkin-0.2.7}/examples/modules/__init__.py +0 -0
  24. {digitalkin-0.2.5 → digitalkin-0.2.7}/examples/modules/storage_module.py +0 -0
  25. {digitalkin-0.2.5 → digitalkin-0.2.7}/setup.cfg +0 -0
  26. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/__init__.py +0 -0
  27. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/grpc_servers/__init__.py +0 -0
  28. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/grpc_servers/_base_server.py +0 -0
  29. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/grpc_servers/module_server.py +0 -0
  30. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/grpc_servers/registry_server.py +0 -0
  31. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/grpc_servers/registry_servicer.py +0 -0
  32. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/grpc_servers/utils/exceptions.py +0 -0
  33. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/grpc_servers/utils/factory.py +0 -0
  34. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/grpc_servers/utils/grpc_client_wrapper.py +0 -0
  35. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/grpc_servers/utils/models.py +0 -0
  36. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/grpc_servers/utils/types.py +0 -0
  37. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/logger.py +0 -0
  38. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/models/__init__.py +0 -0
  39. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/models/module/__init__.py +0 -0
  40. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/models/module/module.py +0 -0
  41. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/models/module/module_types.py +0 -0
  42. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/models/services/__init__.py +0 -0
  43. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/models/services/cost.py +0 -0
  44. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/models/services/storage.py +0 -0
  45. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/modules/__init__.py +0 -0
  46. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/modules/archetype_module.py +0 -0
  47. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/modules/job_manager.py +0 -0
  48. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/modules/tool_module.py +0 -0
  49. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/modules/trigger_module.py +0 -0
  50. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/py.typed +0 -0
  51. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/services/__init__.py +0 -0
  52. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/services/agent/__init__.py +0 -0
  53. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/services/agent/agent_strategy.py +0 -0
  54. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/services/agent/default_agent.py +0 -0
  55. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/services/base_strategy.py +0 -0
  56. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/services/cost/__init__.py +0 -0
  57. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/services/cost/cost_strategy.py +0 -0
  58. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/services/cost/default_cost.py +0 -0
  59. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/services/cost/grpc_cost.py +0 -0
  60. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/services/filesystem/__init__.py +0 -0
  61. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/services/filesystem/default_filesystem.py +0 -0
  62. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/services/filesystem/filesystem_strategy.py +0 -0
  63. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/services/identity/__init__.py +0 -0
  64. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/services/identity/default_identity.py +0 -0
  65. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/services/identity/identity_strategy.py +0 -0
  66. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/services/registry/__init__.py +0 -0
  67. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/services/registry/default_registry.py +0 -0
  68. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/services/registry/registry_strategy.py +0 -0
  69. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/services/services_config.py +0 -0
  70. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/services/services_models.py +0 -0
  71. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/services/setup/default_setup.py +0 -0
  72. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/services/setup/grpc_setup.py +0 -0
  73. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/services/setup/setup_strategy.py +0 -0
  74. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/services/snapshot/__init__.py +0 -0
  75. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/services/snapshot/default_snapshot.py +0 -0
  76. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/services/snapshot/snapshot_strategy.py +0 -0
  77. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/services/storage/__init__.py +0 -0
  78. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/services/storage/default_storage.py +0 -0
  79. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/services/storage/storage_strategy.py +0 -0
  80. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/utils/__init__.py +0 -0
  81. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin/utils/arg_parser.py +0 -0
  82. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin.egg-info/dependency_links.txt +0 -0
  83. {digitalkin-0.2.5 → digitalkin-0.2.7}/src/digitalkin.egg-info/requires.txt +0 -0
  84. {digitalkin-0.2.5 → digitalkin-0.2.7}/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.5
3
+ Version: 0.2.7
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
@@ -8,7 +8,7 @@ import grpc
8
8
  import openai
9
9
  from pydantic import BaseModel
10
10
 
11
- from digitalkin.grpc_servers.utils.models import SecurityMode, ServerConfig, ServerMode
11
+ from digitalkin.grpc_servers.utils.models import SecurityMode, ClientConfig, ServerMode
12
12
  from digitalkin.modules._base_module import BaseModule
13
13
  from digitalkin.services.setup.setup_strategy import SetupData
14
14
 
@@ -45,12 +45,11 @@ class OpenAIToolSecret(BaseModel):
45
45
  """Secret model defining module configuration parameters."""
46
46
 
47
47
 
48
- server_config = ServerConfig(
48
+ client_config = ClientConfig(
49
49
  host="[::]",
50
50
  port=50151,
51
51
  mode=ServerMode.ASYNC,
52
52
  security=SecurityMode.INSECURE,
53
- max_workers=10,
54
53
  credentials=None,
55
54
  )
56
55
 
@@ -81,11 +80,11 @@ class OpenAIToolModule(BaseModule[OpenAIToolInput, OpenAIToolOutput, OpenAIToolS
81
80
  services_config_params = {
82
81
  "storage": {
83
82
  "config": {"setups": OpenAIToolSetup},
84
- "server_config": server_config,
83
+ "client_config": client_config,
85
84
  },
86
85
  "filesystem": {
87
86
  "config": {},
88
- "server_config": server_config,
87
+ "client_config": client_config,
89
88
  },
90
89
  }
91
90
 
@@ -6,7 +6,7 @@ from typing import Any, ClassVar
6
6
 
7
7
  from pydantic import BaseModel
8
8
 
9
- from digitalkin.grpc_servers.utils.models import SecurityMode, ServerConfig, ServerMode
9
+ from digitalkin.grpc_servers.utils.models import SecurityMode, ClientConfig, ServerMode
10
10
  from digitalkin.modules._base_module import BaseModule
11
11
  from digitalkin.services.setup.setup_strategy import SetupData
12
12
  from digitalkin.services.storage.storage_strategy import DataType, StorageRecord
@@ -54,12 +54,11 @@ class TextTransformStorage(BaseModel):
54
54
  ended: bool = False
55
55
 
56
56
 
57
- server_config = ServerConfig(
57
+ client_config = ClientConfig(
58
58
  host="[::]",
59
59
  port=50151,
60
60
  mode=ServerMode.ASYNC,
61
61
  security=SecurityMode.INSECURE,
62
- max_workers=10,
63
62
  credentials=None,
64
63
  )
65
64
 
@@ -91,8 +90,12 @@ class TextTransformModule(BaseModule[TextTransformInput, TextTransformOutput, Te
91
90
  services_config_params = {
92
91
  "storage": {
93
92
  "config": {"monitor": TextTransformStorage, "setups": TextTransformStorage},
94
- "server_config": server_config,
95
- }
93
+ "client_config": client_config,
94
+ },
95
+ "filesystem": {
96
+ "config": {},
97
+ "client_config": client_config,
98
+ },
96
99
  }
97
100
 
98
101
  async def initialize(self, setup_data: SetupData) -> None:
@@ -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.5"
15
+ version = "0.2.7"
16
16
 
17
17
  classifiers = [
18
18
  "Development Status :: 3 - Alpha",
@@ -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.5"
8
+ __version__ = "0.2.7"
@@ -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)
@@ -10,7 +10,7 @@ from digitalkin_proto.digitalkin.filesystem.v2.filesystem_pb2 import FileType as
10
10
 
11
11
  from digitalkin.grpc_servers.utils.exceptions import ServerError
12
12
  from digitalkin.grpc_servers.utils.grpc_client_wrapper import GrpcClientWrapper
13
- from digitalkin.grpc_servers.utils.models import ServerConfig
13
+ from digitalkin.grpc_servers.utils.models import ClientConfig
14
14
  from digitalkin.services.filesystem.filesystem_strategy import (
15
15
  FilesystemData,
16
16
  FilesystemServiceError,
@@ -28,7 +28,7 @@ class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper):
28
28
  self,
29
29
  mission_id: str,
30
30
  config: dict[str, str],
31
- server_config: ServerConfig,
31
+ client_config: ClientConfig,
32
32
  **kwargs, # noqa: ANN003, ARG002
33
33
  ) -> None:
34
34
  """Initialize the default filesystem strategy.
@@ -36,11 +36,11 @@ class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper):
36
36
  Args:
37
37
  mission_id: The ID of the mission this strategy is associated with
38
38
  config: A dictionary mapping names to Pydantic model classes
39
- server_config: The server configuration object
39
+ client_config: The client configuration object
40
40
  kwargs: other optional arguments to pass to the parent class constructor
41
41
  """
42
42
  super().__init__(mission_id, config)
43
- channel = self._init_channel(server_config)
43
+ channel = self._init_channel(client_config)
44
44
  self.stub = filesystem_service_pb2_grpc.FilesystemServiceStub(channel)
45
45
  logger.info("Channel client 'Filesystem' initialized succesfully")
46
46
 
@@ -8,8 +8,8 @@ from google.protobuf.struct_pb2 import Struct
8
8
  from pydantic import BaseModel
9
9
 
10
10
  from digitalkin.grpc_servers.utils.grpc_client_wrapper import GrpcClientWrapper
11
- from digitalkin.grpc_servers.utils.models import ServerConfig
12
- from digitalkin.services.storage.storage_strategy import StorageRecord, StorageServiceError, StorageStrategy
11
+ from digitalkin.grpc_servers.utils.models import ClientConfig
12
+ from digitalkin.services.storage.storage_strategy import DataType, StorageRecord, StorageServiceError, StorageStrategy
13
13
 
14
14
  logger = logging.getLogger(__name__)
15
15
 
@@ -21,13 +21,13 @@ class GrpcStorage(StorageStrategy, GrpcClientWrapper):
21
21
  self,
22
22
  mission_id: str,
23
23
  config: dict[str, type[BaseModel]],
24
- server_config: ServerConfig,
24
+ client_config: ClientConfig,
25
25
  **kwargs, # noqa: ANN003, ARG002
26
26
  ) -> None:
27
27
  """Initialize the storage."""
28
28
  super().__init__(mission_id=mission_id, config=config)
29
29
 
30
- channel = self._init_channel(server_config)
30
+ channel = self._init_channel(client_config)
31
31
  self.stub = storage_service_pb2_grpc.StorageServiceStub(channel)
32
32
  logger.info("Channel client 'storage' initialized succesfully")
33
33
 
@@ -98,7 +98,7 @@ class GrpcStorage(StorageStrategy, GrpcClientWrapper):
98
98
 
99
99
  request = data_pb2.ModifyRecordRequest(data=data_struct, mission_id=self.mission_id, name=name)
100
100
  response: data_pb2.ModifyRecordResponse = self.exec_grpc_query("ModifyRecord", request)
101
- return StorageRecord(**json_format.MessageToDict(response.stored_data))
101
+ return self._build_record_from_proto(response.stored_data, name)
102
102
  except Exception:
103
103
  msg = f"Error while modifing record {name}"
104
104
  logger.exception(msg)
@@ -121,3 +121,36 @@ class GrpcStorage(StorageStrategy, GrpcClientWrapper):
121
121
  logger.exception(msg)
122
122
  return False
123
123
  return True
124
+
125
+ def _build_record_from_proto(self, stored_data: data_pb2.StorageRecord, default_name: str) -> StorageRecord:
126
+ """Helper to construire un StorageRecord complet à partir du message gRPC.
127
+
128
+ Args:
129
+ stored_data: Le message gRPC contenant les données stockées.
130
+ default_name: Le nom par défaut à utiliser si le nom n'est pas présent dans les données.
131
+
132
+ Returns:
133
+ StorageRecord: Un objet StorageRecord construit à partir des données stockées.
134
+ """
135
+ # Converti en dict, avec tous les champs même s'ils sont absents
136
+ raw = json_format.MessageToDict(
137
+ stored_data,
138
+ preserving_proto_field_name=True,
139
+ always_print_fields_with_no_presence=True,
140
+ )
141
+ # On récupère ou on complète les champs obligatoires
142
+ name = raw.get("name", default_name)
143
+ dtype = raw.get("data_type", DataType.OUTPUT.name)
144
+ payload = raw.get("data", {})
145
+
146
+ # Valide le modèle pydantic pour le champ `data`
147
+ validated = self._validate_data(name, payload)
148
+
149
+ return StorageRecord(
150
+ mission_id=self.mission_id,
151
+ name=name,
152
+ data_type=DataType[dtype],
153
+ data=validated,
154
+ creation_date=raw.get("creation_date"),
155
+ update_date=raw.get("update_date"),
156
+ )
@@ -0,0 +1,75 @@
1
+ """LLM format schema for Pydantic models.
2
+
3
+ This module provides functionality to generate JSON schemas for Pydantic models ready for LLMs.
4
+ """
5
+
6
+ import copy
7
+ from typing import Any
8
+
9
+ from pydantic import BaseModel
10
+ from pydantic.json_schema import GenerateJsonSchema, JsonSchemaValue
11
+
12
+
13
+ class CustomOrderSchema(GenerateJsonSchema):
14
+ """Custom schema generator to sort keys in a specific order."""
15
+
16
+ def sort(self, value: JsonSchemaValue, parent_key: str | None = None) -> JsonSchemaValue: # noqa: ARG002
17
+ """Sort the keys of the schema in a specific order.
18
+
19
+ Args:
20
+ value: The schema value to sort.
21
+ parent_key: The parent key of the schema value.
22
+
23
+ Returns:
24
+ The sorted schema value.
25
+ """
26
+ if isinstance(value, dict):
27
+ # Define your preferred order
28
+ preferred = ["title", "description", "type", "examples", "properties"]
29
+ # Collect all keys, putting preferred ones first
30
+ keys = preferred + [k for k in value if k not in preferred]
31
+ # Recurse for each value
32
+ return {k: self.sort(value[k], k) for k in keys if k in value}
33
+ if isinstance(value, list):
34
+ return [self.sort(v) for v in value]
35
+ return value
36
+
37
+
38
+ def inline_refs(schema: dict) -> dict:
39
+ """Recursively resolve and inline all $ref in the schema.
40
+
41
+ Args:
42
+ schema: The JSON schema to inline.
43
+
44
+ Returns:
45
+ The inlined JSON schema.
46
+ """
47
+ schema = copy.deepcopy(schema)
48
+ defs = schema.pop("$defs", {})
49
+
50
+ def _resolve(obj: Any) -> Any: # noqa: ANN401
51
+ if isinstance(obj, dict):
52
+ if "$ref" in obj:
53
+ ref = obj["$ref"]
54
+ if ref.startswith("#/$defs/"):
55
+ key = ref.split("/")[-1]
56
+ return _resolve(defs[key])
57
+ return {k: _resolve(v) for k, v in obj.items()}
58
+ if isinstance(obj, list):
59
+ return [_resolve(item) for item in obj]
60
+ return obj
61
+
62
+ return _resolve(schema)
63
+
64
+
65
+ def llm_ready_schema(model: type[BaseModel]) -> dict:
66
+ """Convert a Pydantic model to a JSON schema ready for LLMs.
67
+
68
+ Args:
69
+ model: The Pydantic model to convert.
70
+
71
+ Returns:
72
+ The JSON schema as a dictionary.
73
+ """
74
+ schema = model.model_json_schema(schema_generator=CustomOrderSchema)
75
+ return inline_refs(schema)
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: digitalkin
3
- Version: 0.2.5
3
+ Version: 0.2.7
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
@@ -78,4 +78,5 @@ src/digitalkin/services/storage/default_storage.py
78
78
  src/digitalkin/services/storage/grpc_storage.py
79
79
  src/digitalkin/services/storage/storage_strategy.py
80
80
  src/digitalkin/utils/__init__.py
81
- src/digitalkin/utils/arg_parser.py
81
+ src/digitalkin/utils/arg_parser.py
82
+ src/digitalkin/utils/llm_ready_schema.py
File without changes
File without changes
File without changes