digitalkin 0.1.1__py3-none-any.whl → 0.2.0__py3-none-any.whl

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 (77) hide show
  1. base_server/__init__.py +1 -0
  2. base_server/mock/__init__.py +5 -0
  3. base_server/mock/mock_pb2.py +39 -0
  4. base_server/mock/mock_pb2_grpc.py +102 -0
  5. base_server/server_async_insecure.py +124 -0
  6. base_server/server_async_secure.py +142 -0
  7. base_server/server_sync_insecure.py +102 -0
  8. base_server/server_sync_secure.py +121 -0
  9. digitalkin/__init__.py +1 -11
  10. digitalkin/__version__.py +1 -4
  11. digitalkin/{grpc → grpc_servers}/__init__.py +1 -13
  12. digitalkin/{grpc → grpc_servers}/_base_server.py +3 -3
  13. digitalkin/{grpc → grpc_servers}/module_server.py +30 -12
  14. digitalkin/{grpc → grpc_servers}/module_servicer.py +30 -14
  15. digitalkin/{grpc → grpc_servers}/registry_server.py +6 -4
  16. digitalkin/{grpc → grpc_servers}/registry_servicer.py +8 -2
  17. digitalkin/{grpc → grpc_servers}/utils/factory.py +6 -4
  18. digitalkin/grpc_servers/utils/grpc_client_wrapper.py +68 -0
  19. digitalkin/{grpc → grpc_servers}/utils/models.py +1 -1
  20. digitalkin/models/__init__.py +1 -4
  21. digitalkin/models/module/__init__.py +8 -2
  22. digitalkin/models/module/module_types.py +10 -0
  23. digitalkin/models/services/__init__.py +0 -5
  24. digitalkin/modules/__init__.py +3 -3
  25. digitalkin/modules/_base_module.py +64 -27
  26. digitalkin/modules/archetype_module.py +2 -6
  27. digitalkin/modules/job_manager.py +46 -28
  28. digitalkin/modules/tool_module.py +3 -7
  29. digitalkin/modules/trigger_module.py +2 -7
  30. digitalkin/services/__init__.py +7 -9
  31. digitalkin/services/agent/__init__.py +2 -2
  32. digitalkin/services/agent/agent_strategy.py +3 -6
  33. digitalkin/services/agent/default_agent.py +1 -4
  34. digitalkin/services/base_strategy.py +18 -0
  35. digitalkin/services/cost/__init__.py +4 -3
  36. digitalkin/services/cost/cost_strategy.py +35 -5
  37. digitalkin/services/cost/default_cost.py +22 -5
  38. digitalkin/services/cost/grpc_cost.py +81 -0
  39. digitalkin/services/filesystem/__init__.py +4 -3
  40. digitalkin/services/filesystem/default_filesystem.py +197 -17
  41. digitalkin/services/filesystem/filesystem_strategy.py +54 -15
  42. digitalkin/services/filesystem/grpc_filesystem.py +209 -0
  43. digitalkin/services/identity/__init__.py +2 -2
  44. digitalkin/services/identity/default_identity.py +1 -1
  45. digitalkin/services/identity/identity_strategy.py +3 -1
  46. digitalkin/services/registry/__init__.py +2 -2
  47. digitalkin/services/registry/default_registry.py +1 -4
  48. digitalkin/services/registry/registry_strategy.py +3 -6
  49. digitalkin/services/services_config.py +176 -0
  50. digitalkin/services/services_models.py +61 -0
  51. digitalkin/services/setup/default_setup.py +222 -0
  52. digitalkin/services/setup/grpc_setup.py +307 -0
  53. digitalkin/services/setup/setup_strategy.py +145 -0
  54. digitalkin/services/snapshot/__init__.py +2 -2
  55. digitalkin/services/snapshot/default_snapshot.py +1 -1
  56. digitalkin/services/snapshot/snapshot_strategy.py +3 -4
  57. digitalkin/services/storage/__init__.py +4 -3
  58. digitalkin/services/storage/default_storage.py +184 -57
  59. digitalkin/services/storage/grpc_storage.py +76 -170
  60. digitalkin/services/storage/storage_strategy.py +195 -24
  61. digitalkin/utils/arg_parser.py +16 -17
  62. {digitalkin-0.1.1.dist-info → digitalkin-0.2.0.dist-info}/METADATA +8 -7
  63. digitalkin-0.2.0.dist-info/RECORD +78 -0
  64. {digitalkin-0.1.1.dist-info → digitalkin-0.2.0.dist-info}/WHEEL +1 -1
  65. digitalkin-0.2.0.dist-info/top_level.txt +3 -0
  66. modules/__init__.py +0 -0
  67. modules/minimal_llm_module.py +162 -0
  68. modules/storage_module.py +187 -0
  69. modules/text_transform_module.py +201 -0
  70. digitalkin/services/default_service.py +0 -13
  71. digitalkin/services/development_service.py +0 -10
  72. digitalkin/services/service_provider.py +0 -27
  73. digitalkin-0.1.1.dist-info/RECORD +0 -59
  74. digitalkin-0.1.1.dist-info/top_level.txt +0 -1
  75. /digitalkin/{grpc → grpc_servers}/utils/exceptions.py +0 -0
  76. /digitalkin/{grpc → grpc_servers}/utils/types.py +0 -0
  77. {digitalkin-0.1.1.dist-info → digitalkin-0.2.0.dist-info}/licenses/LICENSE +0 -0
@@ -1,31 +1,70 @@
1
1
  """This module contains the abstract base class for filesystem strategies."""
2
2
 
3
3
  from abc import ABC, abstractmethod
4
- from typing import Any
4
+ from enum import Enum, auto
5
5
 
6
+ from pydantic import BaseModel, Field
6
7
 
7
- class FilesystemStrategy(ABC):
8
- """Abstract base class for file system strategies."""
8
+ from digitalkin.services.base_strategy import BaseStrategy
9
9
 
10
- def __init__(self) -> None:
11
- """Initialize the file system strategy."""
10
+
11
+ class FilesystemServiceError(Exception):
12
+ """Base exception for Setup service errors."""
13
+
14
+
15
+ class FileType(Enum):
16
+ """Enum defining the types of data that can be stored."""
17
+
18
+ DOCUMENT = auto()
19
+ IMAGE = auto()
20
+ VIDEO = auto()
21
+ AUDIO = auto()
22
+ ARCHIVE = auto()
23
+ OTHER = auto()
24
+
25
+
26
+ class FilesystemData(BaseModel):
27
+ """Data model for filesystem operations."""
28
+
29
+ kin_context: str = Field(description="The context of the file in the filesystem")
30
+ name: str = Field(description="The name of the file")
31
+ file_type: FileType = Field(default=FileType.DOCUMENT, description="The type of data stored")
32
+ url: str = Field(description="The URL of the file in the filesystem")
33
+
34
+
35
+ class FilesystemStrategy(BaseStrategy, ABC):
36
+ """Abstract base class for filesystem strategies."""
37
+
38
+ def __init__(self, mission_id: str, config: dict[str, str]) -> None:
39
+ """Initialize the strategy.
40
+
41
+ Args:
42
+ mission_id: The ID of the mission this strategy is associated with
43
+ config: configuration dictionary for the filesystem strategy
44
+ """
45
+ super().__init__(mission_id)
46
+ self.config: dict[str, str] = config
47
+
48
+ @abstractmethod
49
+ def upload(self, content: bytes, name: str, file_type: FileType) -> FilesystemData:
50
+ """Create a new file in the filesystem."""
12
51
 
13
52
  @abstractmethod
14
- def create(self, data: dict[str, Any]) -> str:
15
- """Create a new file in the file system."""
53
+ def get(self, name: str) -> FilesystemData:
54
+ """Get file from the filesystem."""
16
55
 
17
56
  @abstractmethod
18
- def get(self, data: dict[str, Any]) -> list[dict[str, Any]]:
19
- """Get files from the file system."""
57
+ def get_batch(self, names: list[str]) -> dict[str, FilesystemData | None]:
58
+ """Get files from the filesystem."""
20
59
 
21
60
  @abstractmethod
22
- def update(self, data: dict[str, Any]) -> int:
23
- """Update files in the file system."""
61
+ def get_all(self) -> list[FilesystemData]:
62
+ """Get all files from the filesystem."""
24
63
 
25
64
  @abstractmethod
26
- def delete(self, data: dict[str, Any]) -> int:
27
- """Delete files from the file system."""
65
+ def update(self, name: str, content: bytes, file_type: FileType) -> FilesystemData:
66
+ """Update files in the filesystem."""
28
67
 
29
68
  @abstractmethod
30
- def get_all(self) -> list[dict[str, Any]]:
31
- """Get all files from the file system."""
69
+ def delete(self, name: str) -> bool:
70
+ """Delete file from the filesystem."""
@@ -0,0 +1,209 @@
1
+ """Grpc filesystem."""
2
+
3
+ import logging
4
+ from collections.abc import Generator
5
+ from contextlib import contextmanager
6
+ from typing import Any
7
+
8
+ from digitalkin_proto.digitalkin.filesystem.v2 import filesystem_pb2, filesystem_service_pb2_grpc
9
+ from digitalkin_proto.digitalkin.filesystem.v2.filesystem_pb2 import FileType as FileTypeProto
10
+
11
+ from digitalkin.grpc_servers.utils.exceptions import ServerError
12
+ from digitalkin.grpc_servers.utils.grpc_client_wrapper import GrpcClientWrapper
13
+ from digitalkin.grpc_servers.utils.models import ServerConfig
14
+ from digitalkin.services.filesystem.filesystem_strategy import (
15
+ FilesystemData,
16
+ FilesystemServiceError,
17
+ FilesystemStrategy,
18
+ FileType,
19
+ )
20
+
21
+ logger = logging.getLogger(__name__)
22
+
23
+
24
+ class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper):
25
+ """Default state filesystem strategy."""
26
+
27
+ def __init__(
28
+ self,
29
+ mission_id: str,
30
+ config: dict[str, str],
31
+ server_config: ServerConfig,
32
+ **kwargs, # noqa: ANN003, ARG002
33
+ ) -> None:
34
+ """Initialize the default filesystem strategy.
35
+
36
+ Args:
37
+ mission_id: The ID of the mission this strategy is associated with
38
+ config: A dictionary mapping names to Pydantic model classes
39
+ server_config: The server configuration object
40
+ kwargs: other optional arguments to pass to the parent class constructor
41
+ """
42
+ super().__init__(mission_id, config)
43
+ channel = self._init_channel(server_config)
44
+ self.stub = filesystem_service_pb2_grpc.FilesystemServiceStub(channel)
45
+ logger.info("Channel client 'Filesystem' initialized succesfully")
46
+
47
+ @staticmethod
48
+ @contextmanager
49
+ def _handle_grpc_errors(operation: str) -> Generator[Any, Any, Any]:
50
+ """Context manager for consistent gRPC error handling.
51
+
52
+ Yields:
53
+ Allow error handling in context.
54
+
55
+ Args:
56
+ operation: Description of the operation being performed.
57
+
58
+ Raises:
59
+ ValueError: Error wiht the model validation.
60
+ ServerError: from gRPC Client.
61
+ FilesystemServiceError: Filesystem service internal.
62
+ """
63
+ try:
64
+ yield
65
+ except ServerError as e:
66
+ msg = f"gRPC {operation} failed: {e}"
67
+ logger.exception(msg)
68
+ raise ServerError(msg) from e
69
+ except Exception as e:
70
+ msg = f"Unexpected error in {operation}"
71
+ logger.exception(msg)
72
+ raise FilesystemServiceError(msg) from e
73
+
74
+ def upload(self, content: bytes, name: str, file_type: FileType) -> FilesystemData:
75
+ """Create a new file in the file system.
76
+
77
+ Args:
78
+ content: The content of the file to be uploaded
79
+ name: The name of the file to be created
80
+ file_type: The type of data being uploaded
81
+
82
+ Returns:
83
+ FilesystemData: Metadata about the uploaded file
84
+
85
+ Raises:
86
+ ValueError: If the file already exists
87
+ """
88
+ with GrpcFilesystem._handle_grpc_errors("UploadFile"):
89
+ request = filesystem_pb2.UploadFileRequest(
90
+ kin_context=self.mission_id,
91
+ name=name,
92
+ file_type=file_type.name,
93
+ content=content,
94
+ )
95
+ response: filesystem_pb2.UploadFileResponse = self.exec_grpc_query("UploadFile", request)
96
+ return FilesystemData(
97
+ kin_context=response.file.kin_context,
98
+ name=response.file.name,
99
+ file_type=FileType[FileTypeProto.Name(response.file.file_type)],
100
+ url=response.file.url,
101
+ )
102
+
103
+ def get(self, name: str) -> FilesystemData:
104
+ """Get file from the filesystem.
105
+
106
+ Args:
107
+ name: The name of the file to be retrieved
108
+
109
+ Returns:
110
+ FilesystemData: Metadata about the retrieved file
111
+ """
112
+ with GrpcFilesystem._handle_grpc_errors("GetFileByName"):
113
+ request = filesystem_pb2.GetFileByNameRequest(name=name)
114
+ response: filesystem_pb2.GetFileByNameResponse = self.exec_grpc_query("GetFileByName", request)
115
+ return FilesystemData(
116
+ kin_context=response.file.kin_context,
117
+ name=response.file.name,
118
+ file_type=FileType[FileTypeProto.Name(response.file.file_type)],
119
+ url=response.file.url,
120
+ )
121
+
122
+ def update(self, name: str, content: bytes, file_type: FileType) -> FilesystemData:
123
+ """Update files in the filesystem.
124
+
125
+ Args:
126
+ name: The name of the file to be updated
127
+ content: The new content of the file
128
+ file_type: The type of data being uploaded
129
+
130
+ Returns:
131
+ FilesystemData: Metadata about the updated file
132
+ """
133
+ with GrpcFilesystem._handle_grpc_errors("UpdateFile"):
134
+ request = filesystem_pb2.UpdateFileRequest(
135
+ kin_context=self.mission_id,
136
+ name=name,
137
+ file_type=file_type.name,
138
+ content=content,
139
+ )
140
+ response: filesystem_pb2.UpdateFileResponse = self.exec_grpc_query("UpdateFile", request)
141
+ return FilesystemData(
142
+ kin_context=response.file.kin_context,
143
+ name=response.file.name,
144
+ file_type=FileType[FileTypeProto.Name(response.file.file_type)],
145
+ url=response.file.url,
146
+ )
147
+
148
+ def delete(self, name: str) -> bool:
149
+ """Delete files from the filesystem.
150
+
151
+ Args:
152
+ name: The name of the file to be deleted
153
+
154
+ Returns:
155
+ int: 1 if the file was deleted successfully, 0 if it didn't exist, None on error
156
+ """
157
+ with GrpcFilesystem._handle_grpc_errors("DeleteFile"):
158
+ request = filesystem_pb2.DeleteFileRequest(name=name)
159
+ _: filesystem_pb2.DeleteFileResponse = self.exec_grpc_query("DeleteFile", request)
160
+ return True
161
+
162
+ def get_all(self) -> list[FilesystemData]:
163
+ """Get all files from the filesystem.
164
+
165
+ Returns:
166
+ list[FilesystemData]: A list of all files in the filesystem
167
+ """
168
+ with GrpcFilesystem._handle_grpc_errors("GetFilesByKinContext"):
169
+ request = filesystem_pb2.GetFilesByKinContextRequest(kin_context=self.mission_id)
170
+ response: filesystem_pb2.GetFilesByKinContextResponse = self.exec_grpc_query(
171
+ "GetFilesByKinContext", request
172
+ )
173
+ return [
174
+ FilesystemData(
175
+ kin_context=file.kin_context,
176
+ name=file.name,
177
+ file_type=FileType[FileTypeProto.Name(file.file_type)],
178
+ url=file.url,
179
+ )
180
+ for file in response.files
181
+ ]
182
+
183
+ def get_batch(self, names: list[str]) -> dict[str, FilesystemData | None]:
184
+ """Get files from the filesystem.
185
+
186
+ Args:
187
+ names: The names of the files to be retrieved
188
+
189
+ Returns:
190
+ list[FilesystemData]: A list of metadata about the retrieved files
191
+ """
192
+ with GrpcFilesystem._handle_grpc_errors("GetFilesByNames"):
193
+ request = filesystem_pb2.GetFilesByNamesRequest(names=names)
194
+ response: filesystem_pb2.GetFilesByNamesResponse = self.exec_grpc_query("GetFilesByNames", request)
195
+ result: dict[str, FilesystemData | None] = {}
196
+ for name, file_result in response.files.items():
197
+ which_field = file_result.WhichOneof("result")
198
+ if which_field == "file":
199
+ result[name] = FilesystemData(
200
+ kin_context=file_result.file.kin_context,
201
+ name=file_result.file.name,
202
+ file_type=FileType[FileTypeProto.Name(file_result.file.file_type)],
203
+ url=file_result.file.url,
204
+ )
205
+ elif which_field == "error":
206
+ # Handle error case
207
+ result[name] = None
208
+ logger.warning("Error retrieving file '%s': %s", name, file_result.error)
209
+ return result
@@ -1,6 +1,6 @@
1
1
  """This module is responsible for handling the identity service."""
2
2
 
3
- from .default_identity import DefaultIdentity
4
- from .identity_strategy import IdentityStrategy
3
+ from digitalkin.services.identity.default_identity import DefaultIdentity
4
+ from digitalkin.services.identity.identity_strategy import IdentityStrategy
5
5
 
6
6
  __all__ = ["DefaultIdentity", "IdentityStrategy"]
@@ -1,6 +1,6 @@
1
1
  """Default identity."""
2
2
 
3
- from .identity_strategy import IdentityStrategy
3
+ from digitalkin.services.identity.identity_strategy import IdentityStrategy
4
4
 
5
5
 
6
6
  class DefaultIdentity(IdentityStrategy):
@@ -2,8 +2,10 @@
2
2
 
3
3
  from abc import ABC, abstractmethod
4
4
 
5
+ from digitalkin.services.base_strategy import BaseStrategy
5
6
 
6
- class IdentityStrategy(ABC):
7
+
8
+ class IdentityStrategy(BaseStrategy, ABC):
7
9
  """IdentityStrategy is the abstract base class for all identity strategies."""
8
10
 
9
11
  @abstractmethod
@@ -1,6 +1,6 @@
1
1
  """This module is responsible for handling the registry service."""
2
2
 
3
- from .default_registry import DefaultRegistry
4
- from .registry_strategy import RegistryStrategy
3
+ from digitalkin.services.registry.default_registry import DefaultRegistry
4
+ from digitalkin.services.registry.registry_strategy import RegistryStrategy
5
5
 
6
6
  __all__ = ["DefaultRegistry", "RegistryStrategy"]
@@ -1,13 +1,10 @@
1
1
  """Default registry."""
2
2
 
3
- from .registry_strategy import RegistryStrategy
3
+ from digitalkin.services.registry.registry_strategy import RegistryStrategy
4
4
 
5
5
 
6
6
  class DefaultRegistry(RegistryStrategy):
7
7
  """Default registry strategy."""
8
8
 
9
- def __init__(self) -> None:
10
- """Initialize the registry strategy."""
11
-
12
9
  def get_by_id(self, module_id: str) -> None:
13
10
  """Get services from the registry."""
@@ -2,14 +2,11 @@
2
2
 
3
3
  from abc import ABC, abstractmethod
4
4
 
5
+ from digitalkin.services.base_strategy import BaseStrategy
5
6
 
6
- class RegistryStrategy(ABC):
7
- """Abstract base class for registry strategies."""
8
7
 
9
- @abstractmethod
10
- def __init__(self) -> None:
11
- """Initialize the registry strategy."""
12
- raise NotImplementedError
8
+ class RegistryStrategy(BaseStrategy, ABC):
9
+ """Abstract base class for registry strategies."""
13
10
 
14
11
  @abstractmethod
15
12
  def get_by_id(self, module_id: str) -> None:
@@ -0,0 +1,176 @@
1
+ """Service Provider definitions."""
2
+
3
+ from typing import Any, ClassVar
4
+
5
+ from pydantic import BaseModel, Field, PrivateAttr
6
+
7
+ from digitalkin.services.agent import AgentStrategy, DefaultAgent
8
+ from digitalkin.services.cost import CostStrategy, DefaultCost
9
+ from digitalkin.services.filesystem import DefaultFilesystem, FilesystemStrategy, GrpcFilesystem
10
+ from digitalkin.services.identity import DefaultIdentity, IdentityStrategy
11
+ from digitalkin.services.registry import DefaultRegistry, RegistryStrategy
12
+ from digitalkin.services.services_models import ServicesMode, ServicesStrategy
13
+ from digitalkin.services.snapshot import DefaultSnapshot, SnapshotStrategy
14
+ from digitalkin.services.storage import DefaultStorage, GrpcStorage, StorageStrategy
15
+
16
+
17
+ class ServicesConfig(BaseModel):
18
+ """Service class describing the available services in a Module.
19
+
20
+ This class manages the strategy implementations for various services,
21
+ allowing them to be switched between local and remote modes.
22
+ """
23
+
24
+ # Mode setting for all strategies
25
+ mode: ServicesMode = Field(default=ServicesMode.LOCAL, description="The mode of the services (local or remote)")
26
+
27
+ # Strategy definitions with proper type annotations
28
+ _storage: ServicesStrategy[StorageStrategy] = PrivateAttr(
29
+ default_factory=lambda: ServicesStrategy(local=DefaultStorage, remote=GrpcStorage)
30
+ )
31
+ _config_storage: dict[str, Any | None] = PrivateAttr(default_factory=dict)
32
+ _cost: ServicesStrategy[CostStrategy] = PrivateAttr(
33
+ default_factory=lambda: ServicesStrategy(local=DefaultCost, remote=DefaultCost)
34
+ )
35
+ _config_cost: dict[str, Any | None] = PrivateAttr(default_factory=dict)
36
+ _snapshot: ServicesStrategy[SnapshotStrategy] = PrivateAttr(
37
+ default_factory=lambda: ServicesStrategy(local=DefaultSnapshot, remote=DefaultSnapshot)
38
+ )
39
+ _config_snapshot: dict[str, Any | None] = PrivateAttr(default_factory=dict)
40
+ _registry: ServicesStrategy[RegistryStrategy] = PrivateAttr(
41
+ default_factory=lambda: ServicesStrategy(local=DefaultRegistry, remote=DefaultRegistry)
42
+ )
43
+ _config__registry: dict[str, Any | None] = PrivateAttr(default_factory=dict)
44
+ _filesystem: ServicesStrategy[FilesystemStrategy] = PrivateAttr(
45
+ default_factory=lambda: ServicesStrategy(local=DefaultFilesystem, remote=GrpcFilesystem)
46
+ )
47
+ _config_filesystem: dict[str, Any | None] = PrivateAttr(default_factory=dict)
48
+ _agent: ServicesStrategy[AgentStrategy] = PrivateAttr(
49
+ default_factory=lambda: ServicesStrategy(local=DefaultAgent, remote=DefaultAgent)
50
+ )
51
+ _config_agent: dict[str, Any | None] = PrivateAttr(default_factory=dict)
52
+ _identity: ServicesStrategy[IdentityStrategy] = PrivateAttr(
53
+ default_factory=lambda: ServicesStrategy(local=DefaultIdentity, remote=DefaultIdentity)
54
+ )
55
+ _config_identity: dict[str, Any | None] = PrivateAttr(default_factory=dict)
56
+
57
+ # List of valid strategy names for validation
58
+ _valid_strategy_names: ClassVar[set[str]] = {
59
+ "storage",
60
+ "cost",
61
+ "snapshot",
62
+ "registry",
63
+ "filesystem",
64
+ "agent",
65
+ "identity",
66
+ }
67
+
68
+ def __init__(
69
+ self,
70
+ services_config_strategies: dict[str, ServicesStrategy | None] = {},
71
+ services_config_params: dict[str, dict[str, Any | None] | None] = {},
72
+ mode: ServicesMode = ServicesMode.LOCAL,
73
+ **kwargs: dict[str, Any],
74
+ ) -> None:
75
+ """Initialize the service configuration with optional strategy overrides.
76
+
77
+ Args:
78
+ services_config_strategies: Dictionary mapping service names to strategy implementations
79
+ services_config_params: Dictionary mapping service names to configuration parameters
80
+ mode: The mode of the services (local or remote)
81
+ **kwargs: Additional keyword arguments passed to the parent class constructor
82
+ """
83
+ super().__init__(**kwargs)
84
+ self.mode = mode
85
+ # Apply any strategy overrides
86
+ if services_config_strategies:
87
+ for name, strategy in services_config_strategies.items():
88
+ if strategy is not None and name in self._valid_strategy_names:
89
+ setattr(self, f"_{name}", strategy)
90
+
91
+ for name in self.valid_strategy_names():
92
+ setattr(self, f"_config_{name}", services_config_params.get(name, {}))
93
+
94
+ @classmethod
95
+ def valid_strategy_names(cls) -> set[str]:
96
+ """Get the list of valid strategy names.
97
+
98
+ Returns:
99
+ The set of valid strategy names.
100
+ """
101
+ return cls._valid_strategy_names
102
+
103
+ def get_strategy_config(self, name: str) -> dict[str, Any]:
104
+ """Get the configuration for a specific strategy.
105
+
106
+ Args:
107
+ name: The name of the strategy to retrieve the configuration for
108
+
109
+ Returns:
110
+ The configuration for the specified strategy, or None if not found
111
+ """
112
+ return getattr(self, f"_config_{name}", {})
113
+
114
+ def init_strategy(self, name: str, mission_id: str) -> ServicesStrategy:
115
+ """Initialize a specific strategy.
116
+
117
+ Args:
118
+ name: The name of the strategy to initialize
119
+ mission_id: The ID of the mission this strategy is associated with
120
+
121
+ Returns:
122
+ The initialized strategy instance
123
+
124
+ Raises:
125
+ ValueError: If the strategy is not found
126
+ """
127
+ strategy_type = getattr(self, name, None)
128
+ if strategy_type is None:
129
+ msg = f"Strategy {name} not found in ServicesConfig."
130
+ raise ValueError(msg)
131
+
132
+ # Instantiate the strategy with the mission ID and configuration
133
+ return strategy_type(mission_id, **self.get_strategy_config(name) or {})
134
+
135
+ @property
136
+ def storage(self) -> type[StorageStrategy]:
137
+ """Get the storage service strategy class based on the current mode."""
138
+ return self._storage[self.mode.value]
139
+
140
+ @property
141
+ def cost(self) -> type[CostStrategy]:
142
+ """Get the cost service strategy class based on the current mode."""
143
+ return self._cost[self.mode.value]
144
+
145
+ @property
146
+ def snapshot(self) -> type[SnapshotStrategy]:
147
+ """Get the snapshot service strategy class based on the current mode."""
148
+ return self._snapshot[self.mode.value]
149
+
150
+ @property
151
+ def registry(self) -> type[RegistryStrategy]:
152
+ """Get the registry service strategy class based on the current mode."""
153
+ return self._registry[self.mode.value]
154
+
155
+ @property
156
+ def filesystem(self) -> type[FilesystemStrategy]:
157
+ """Get the filesystem service strategy class based on the current mode."""
158
+ return self._filesystem[self.mode.value]
159
+
160
+ @property
161
+ def agent(self) -> type[AgentStrategy]:
162
+ """Get the agent service strategy class based on the current mode."""
163
+ return self._agent[self.mode.value]
164
+
165
+ @property
166
+ def identity(self) -> type[IdentityStrategy]:
167
+ """Get the identity service strategy class based on the current mode."""
168
+ return self._identity[self.mode.value]
169
+
170
+ def update_mode(self, mode: ServicesMode) -> None:
171
+ """Update the strategy mode.
172
+
173
+ Parameters:
174
+ mode: The new mode to use for all strategies
175
+ """
176
+ self.mode = mode
@@ -0,0 +1,61 @@
1
+ """This module contains the strategy models for the services."""
2
+
3
+ from enum import Enum
4
+ from typing import Generic, TypeVar
5
+
6
+ from pydantic import BaseModel
7
+
8
+ from digitalkin.logger import logger
9
+ from digitalkin.services.agent import AgentStrategy
10
+ from digitalkin.services.cost import CostStrategy
11
+ from digitalkin.services.filesystem import FilesystemStrategy
12
+ from digitalkin.services.identity import IdentityStrategy
13
+ from digitalkin.services.registry import RegistryStrategy
14
+ from digitalkin.services.snapshot import SnapshotStrategy
15
+ from digitalkin.services.storage import StorageStrategy
16
+
17
+ # Define type variables
18
+ T = TypeVar(
19
+ "T",
20
+ bound=AgentStrategy
21
+ | CostStrategy
22
+ | FilesystemStrategy
23
+ | IdentityStrategy
24
+ | RegistryStrategy
25
+ | SnapshotStrategy
26
+ | StorageStrategy,
27
+ )
28
+
29
+
30
+ class ServicesMode(str, Enum):
31
+ """Mode for strategy execution."""
32
+
33
+ LOCAL = "local"
34
+ REMOTE = "remote"
35
+
36
+
37
+ class ServicesStrategy(BaseModel, Generic[T]):
38
+ """Service class describing the available services in a Module with local and remote attributes.
39
+
40
+ Attributes:
41
+ local: type
42
+ remote: type
43
+ """
44
+
45
+ local: type[T]
46
+ remote: type[T]
47
+
48
+ def __getitem__(self, mode: str) -> type[T]:
49
+ """Get the service strategy based on the mode.
50
+
51
+ Args:
52
+ mode (str): The mode to get the strategy for.
53
+
54
+ Returns:
55
+ The strategy based on the mode.
56
+ """
57
+ try:
58
+ return getattr(self, mode)
59
+ except AttributeError:
60
+ logger.exception("Unknown mode: %s, available modes are: %s", mode, ServicesMode.__members__)
61
+ return getattr(self, ServicesMode.LOCAL.value)