digitalkin 0.2.11__py3-none-any.whl → 0.2.13__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 (41) hide show
  1. digitalkin/__version__.py +1 -1
  2. digitalkin/grpc_servers/_base_server.py +15 -17
  3. digitalkin/grpc_servers/module_server.py +9 -10
  4. digitalkin/grpc_servers/module_servicer.py +108 -85
  5. digitalkin/grpc_servers/registry_server.py +3 -6
  6. digitalkin/grpc_servers/registry_servicer.py +18 -19
  7. digitalkin/grpc_servers/utils/grpc_client_wrapper.py +3 -5
  8. digitalkin/logger.py +45 -1
  9. digitalkin/models/module/module.py +1 -0
  10. digitalkin/modules/_base_module.py +47 -6
  11. digitalkin/modules/job_manager/base_job_manager.py +139 -0
  12. digitalkin/modules/job_manager/job_manager_models.py +44 -0
  13. digitalkin/modules/job_manager/single_job_manager.py +218 -0
  14. digitalkin/modules/job_manager/taskiq_broker.py +173 -0
  15. digitalkin/modules/job_manager/taskiq_job_manager.py +213 -0
  16. digitalkin/services/base_strategy.py +3 -1
  17. digitalkin/services/cost/cost_strategy.py +64 -16
  18. digitalkin/services/cost/default_cost.py +95 -12
  19. digitalkin/services/cost/grpc_cost.py +149 -60
  20. digitalkin/services/filesystem/default_filesystem.py +5 -6
  21. digitalkin/services/filesystem/filesystem_strategy.py +3 -2
  22. digitalkin/services/filesystem/grpc_filesystem.py +31 -26
  23. digitalkin/services/services_config.py +6 -5
  24. digitalkin/services/setup/__init__.py +1 -0
  25. digitalkin/services/setup/default_setup.py +10 -12
  26. digitalkin/services/setup/grpc_setup.py +8 -10
  27. digitalkin/services/storage/default_storage.py +13 -6
  28. digitalkin/services/storage/grpc_storage.py +25 -9
  29. digitalkin/services/storage/storage_strategy.py +3 -2
  30. digitalkin/utils/arg_parser.py +5 -48
  31. digitalkin/utils/development_mode_action.py +51 -0
  32. {digitalkin-0.2.11.dist-info → digitalkin-0.2.13.dist-info}/METADATA +43 -12
  33. {digitalkin-0.2.11.dist-info → digitalkin-0.2.13.dist-info}/RECORD +40 -33
  34. {digitalkin-0.2.11.dist-info → digitalkin-0.2.13.dist-info}/WHEEL +1 -1
  35. modules/cpu_intensive_module.py +271 -0
  36. modules/minimal_llm_module.py +200 -56
  37. modules/storage_module.py +5 -6
  38. modules/text_transform_module.py +1 -1
  39. digitalkin/modules/job_manager.py +0 -176
  40. {digitalkin-0.2.11.dist-info → digitalkin-0.2.13.dist-info}/licenses/LICENSE +0 -0
  41. {digitalkin-0.2.11.dist-info → digitalkin-0.2.13.dist-info}/top_level.txt +0 -0
@@ -1,8 +1,8 @@
1
1
  """This module contains the abstract base class for cost strategies."""
2
2
 
3
3
  from abc import ABC, abstractmethod
4
- from enum import Enum, auto
5
- from typing import Any
4
+ from enum import Enum
5
+ from typing import Literal
6
6
 
7
7
  from pydantic import BaseModel
8
8
 
@@ -10,36 +10,84 @@ from digitalkin.services.base_strategy import BaseStrategy
10
10
 
11
11
 
12
12
  class CostType(Enum):
13
- """."""
13
+ """Enum defining the types of costs that can be registered."""
14
14
 
15
- OTHER = auto()
16
- TOKEN_INPUT = auto()
17
- TOKEN_OUTPUT = auto()
18
- API_CALL = auto()
19
- STORAGE = auto()
20
- TIME = auto()
15
+ OTHER = "OTHER"
16
+ TOKEN_INPUT = "TOKEN_INPUT"
17
+ TOKEN_OUTPUT = "TOKEN_OUTPUT"
18
+ API_CALL = "API_CALL"
19
+ STORAGE = "STORAGE"
20
+ TIME = "TIME"
21
+
22
+
23
+ class CostConfig(BaseModel):
24
+ """Pydantic model that defines a cost configuration.
25
+
26
+ :param cost_name: Name of the cost (unique identifier in the service).
27
+ :param cost_type: The type/category of the cost.
28
+ :param description: A short description of the cost.
29
+ :param unit: The unit of measurement (e.g. token, call, MB).
30
+ :param rate: The cost per unit (e.g. dollars per token).
31
+ """
32
+
33
+ cost_name: str
34
+ cost_type: Literal["TOKEN_INPUT", "TOKEN_OUTPUT", "API_CALL", "STORAGE", "TIME", "OTHER"]
35
+ description: str | None = None
36
+ unit: str
37
+ rate: float
21
38
 
22
39
 
23
40
  class CostData(BaseModel):
24
- """."""
41
+ """Data model for cost operations."""
25
42
 
26
43
  cost: float
27
44
  mission_id: str
28
45
  name: str
29
- type: CostType
46
+ cost_type: CostType
30
47
  unit: str
48
+ rate: float
49
+ setup_version_id: str
50
+ quantity: float
51
+
52
+
53
+ class CostServiceError(Exception):
54
+ """Custom exception for CostService errors."""
31
55
 
32
56
 
33
57
  class CostStrategy(BaseStrategy, ABC):
34
58
  """Abstract base class for cost strategies."""
35
59
 
60
+ def __init__(self, mission_id: str, setup_version_id: str, config: dict[str, CostConfig]) -> None:
61
+ """Initialize the strategy.
62
+
63
+ Args:
64
+ mission_id: The ID of the mission this strategy is associated with
65
+ setup_version_id: The ID of the setup version this strategy is associated with
66
+ config: Configuration dictionary for the strategy
67
+ """
68
+ super().__init__(mission_id, setup_version_id)
69
+ self.config = config
70
+
36
71
  @abstractmethod
37
- def add_cost(self, cost_dict: dict[str, Any]) -> str:
72
+ def add(
73
+ self,
74
+ name: str,
75
+ cost_config_name: str,
76
+ quantity: float,
77
+ ) -> None:
38
78
  """Register a new cost."""
39
79
 
40
- def __post_init__(self, *args, **kwargs) -> None: # noqa: ANN002, ANN003
41
- """Allow post init configuration."""
42
-
43
80
  @abstractmethod
44
- def get(self, cost_dict: dict[str, Any]) -> list[CostData]:
81
+ def get(
82
+ self,
83
+ name: str,
84
+ ) -> list[CostData]:
45
85
  """Get a cost."""
86
+
87
+ @abstractmethod
88
+ def get_filtered(
89
+ self,
90
+ names: list[str] | None = None,
91
+ cost_types: list[Literal["TOKEN_INPUT", "TOKEN_OUTPUT", "API_CALL", "STORAGE", "TIME", "OTHER"]] | None = None,
92
+ ) -> list[CostData]:
93
+ """Get filtered costs."""
@@ -1,30 +1,113 @@
1
1
  """Default cost."""
2
2
 
3
- import logging
4
- from typing import Any
3
+ from typing import Literal
5
4
 
6
- from digitalkin.services.cost.cost_strategy import CostData, CostStrategy
7
-
8
- logger = logging.getLogger(__name__)
5
+ from digitalkin.logger import logger
6
+ from digitalkin.services.cost.cost_strategy import (
7
+ CostConfig,
8
+ CostData,
9
+ CostServiceError,
10
+ CostStrategy,
11
+ CostType,
12
+ )
9
13
 
10
14
 
11
15
  class DefaultCost(CostStrategy):
12
16
  """Default cost strategy."""
13
17
 
14
- def add_cost(self, cost_dict: dict[str, Any]) -> str: # noqa: PLR6301
18
+ def __init__(self, mission_id: str, setup_version_id: str, config: dict[str, CostConfig]) -> None:
19
+ """Initialize the strategy.
20
+
21
+ Args:
22
+ mission_id: The ID of the mission this strategy is associated with
23
+ setup_version_id: The ID of the setup version this strategy is associated with
24
+ config: The configuration dictionary for the cost
25
+ """
26
+ super().__init__(mission_id=mission_id, setup_version_id=setup_version_id, config=config)
27
+ self.db: dict[str, list[CostData]] = {}
28
+
29
+ def add(
30
+ self,
31
+ name: str,
32
+ cost_config_name: str,
33
+ quantity: float,
34
+ ) -> None:
15
35
  """Create a new record in the cost database.
16
36
 
37
+ Args:
38
+ name: The name of the cost
39
+ cost_config_name: The name of the cost config
40
+ quantity: The quantity of the cost
41
+
42
+ Raises:
43
+ CostServiceError: If the cost data is invalid or if the cost already exists
44
+ """
45
+ cost_config = self.config.get(cost_config_name)
46
+ if cost_config is None:
47
+ msg = f"Cost config {cost_config_name} not found in the configuration."
48
+ logger.error(msg)
49
+ raise CostServiceError(msg)
50
+ cost_data = CostData.model_validate({
51
+ "name": name,
52
+ "cost": cost_config.rate * quantity,
53
+ "unit": cost_config.unit,
54
+ "cost_type": getattr(CostType, cost_config.cost_type),
55
+ "mission_id": self.mission_id,
56
+ "rate": cost_config.rate,
57
+ "quantity": quantity,
58
+ "setup_version_id": self.setup_version_id,
59
+ })
60
+ if cost_data.mission_id not in self.db:
61
+ self.db[cost_data.mission_id] = []
62
+ if cost_data.name in [cost.name for cost in self.db[cost_data.mission_id]]:
63
+ msg = f"Cost with name {cost_data.name} already exists in mission {cost_data.mission_id}"
64
+ logger.error(msg)
65
+ raise CostServiceError(msg)
66
+ self.db[cost_data.mission_id].append(cost_data)
67
+
68
+ def get(self, name: str) -> list[CostData]:
69
+ """Get a record from the database.
70
+
71
+ Args:
72
+ name: The name of the cost
73
+
17
74
  Returns:
18
- str: The ID of the new record
75
+ list[CostData]: The cost data
76
+
77
+ Raises:
78
+ CostServiceError: If the cost data is invalid or if the cost does not exist
19
79
  """
20
- logger.info("Cost added with cost_dict: %s", cost_dict)
21
- return ""
80
+ if self.mission_id not in self.db:
81
+ msg = f"Mission {self.mission_id} not found in the database."
82
+ logger.warning(msg)
83
+ raise CostServiceError(msg)
84
+
85
+ return [cost for cost in self.db[self.mission_id] if cost.name == name] or []
22
86
 
23
- def get(self, cost_dict: dict[str, Any]) -> list[CostData]: # noqa: PLR6301
87
+ def get_filtered(
88
+ self,
89
+ names: list[str] | None = None,
90
+ cost_types: list[Literal["TOKEN_INPUT", "TOKEN_OUTPUT", "API_CALL", "STORAGE", "TIME", "OTHER"]] | None = None,
91
+ ) -> list[CostData]:
24
92
  """Get records from the database.
25
93
 
94
+ Args:
95
+ names: The names of the costs
96
+ cost_types: The types of the costs
97
+
26
98
  Returns:
27
99
  list[CostData]: The list of records
100
+
101
+ Raises:
102
+ CostServiceError: If the cost data is invalid or if the cost does not exist
28
103
  """
29
- logger.info("Costs querried with cost_dict: %s", cost_dict)
30
- return []
104
+ if self.mission_id not in self.db:
105
+ msg = f"Mission {self.mission_id} not found in the database."
106
+ logger.warning(msg)
107
+ raise CostServiceError(msg)
108
+
109
+ return [
110
+ cost
111
+ for cost in self.db[self.mission_id]
112
+ if (names and cost.name in names) or (cost_types and cost.cost_type in cost_types)
113
+ ]
@@ -1,81 +1,170 @@
1
- """This module implements the default Cost strategy."""
1
+ """This module implements the gRPC Cost strategy."""
2
2
 
3
- import logging
4
- from typing import Any
3
+ from collections.abc import Generator
4
+ from contextlib import contextmanager
5
+ from typing import Any, Literal
5
6
 
6
7
  from digitalkin_proto.digitalkin.cost.v1 import cost_pb2, cost_service_pb2_grpc
7
8
  from google.protobuf import json_format
8
- from pydantic import ValidationError
9
9
 
10
+ from digitalkin.grpc_servers.utils.exceptions import ServerError
10
11
  from digitalkin.grpc_servers.utils.grpc_client_wrapper import GrpcClientWrapper
11
- from digitalkin.grpc_servers.utils.models import ServerConfig
12
- from digitalkin.services.cost.cost_strategy import CostData, CostStrategy
13
-
14
- logger = logging.getLogger(__name__)
12
+ from digitalkin.grpc_servers.utils.models import ClientConfig
13
+ from digitalkin.logger import logger
14
+ from digitalkin.services.cost.cost_strategy import (
15
+ CostConfig,
16
+ CostData,
17
+ CostServiceError,
18
+ CostStrategy,
19
+ CostType,
20
+ )
15
21
 
16
22
 
17
23
  class GrpcCost(CostStrategy, GrpcClientWrapper):
18
24
  """This class implements the default Cost strategy."""
19
25
 
20
- def _get_costs_by_name(self, cost_dict: dict[str, Any]) -> list[CostData]:
21
- request = cost_pb2.GetCostsByNameRequest(
22
- mission_id=cost_dict["mission_id"],
23
- name=cost_dict["name"],
24
- )
25
- response = self.exec_grpc_query("GetCostsByName", request)
26
- return [CostData(**json_format.MessageToDict(cost)) for cost in response.costs]
27
-
28
- def _get_costs_by_mission(self, cost_dict: dict[str, Any]) -> list[CostData]:
29
- request = cost_pb2.GetCostsByMissionRequest(mission_id=cost_dict["mission_id"])
30
- response = self.exec_grpc_query("GetCostsByMission", request)
31
- return [CostData(**json_format.MessageToDict(cost)) for cost in response.costs]
32
-
33
- def _get_costs_by_type(self, cost_dict: dict[str, Any]) -> list[CostData]:
34
- request = cost_pb2.GetCostsByTypeRequest(
35
- mission_id=cost_dict["mission_id"],
36
- type=cost_dict["type"],
37
- )
38
- response = self.exec_grpc_query("GetCostsBytype", request)
39
- return [CostData(**json_format.MessageToDict(cost)) for cost in response.costs]
40
-
41
- def __init__(self, mission_id: str, config: ServerConfig) -> None:
42
- """Initialize the cost."""
43
- super().__init__(mission_id)
44
- channel = self._init_channel(config)
45
- self.stub = cost_service_pb2_grpc.CostServiceStub(channel)
46
- logger.info("Channel client 'Cost' initialized succesfully")
26
+ @staticmethod
27
+ @contextmanager
28
+ def _handle_grpc_errors(operation: str) -> Generator[Any, Any, Any]:
29
+ """Context manager for consistent gRPC error handling.
47
30
 
48
- def add_cost(self, cost_dict: dict[str, Any]) -> str:
49
- """Create a new record in the cost database.
31
+ Yields:
32
+ Allow error handling in context.
50
33
 
51
- Required arguments:
52
- data: Object representation of CostData
34
+ Args:
35
+ operation: Description of the operation being performed.
53
36
 
54
- Returns:
55
- str: The ID of the new record
37
+ Raises:
38
+ ValueError: Error with the model validation.
39
+ ServerError: from gRPC Client.
40
+ CostServiceError: Unexpected error.
56
41
  """
57
42
  try:
58
- valid_data = CostData.model_validate(cost_dict["data"])
59
- except ValidationError:
60
- logger.exception("Validation failed for model StorageData")
61
- return ""
62
- except KeyError:
63
- logger.exception("Missing mandatory 'data' in dict.")
64
- return ""
43
+ yield
44
+ except CostServiceError as e:
45
+ msg = f"CostServiceError in {operation}: {e}"
46
+ logger.exception(msg)
47
+ raise CostServiceError(msg) from e
48
+ except ServerError as e:
49
+ msg = f"gRPC {operation} failed: {e}"
50
+ logger.exception(msg)
51
+ raise ServerError(msg) from e
52
+ except Exception as e:
53
+ msg = f"Unexpected error in {operation}"
54
+ logger.exception(msg)
55
+ raise CostServiceError(msg) from e
56
+
57
+ def __init__(
58
+ self,
59
+ mission_id: str,
60
+ setup_version_id: str,
61
+ config: dict[str, CostConfig],
62
+ client_config: ClientConfig,
63
+ ) -> None:
64
+ """Initialize the cost."""
65
+ super().__init__(mission_id=mission_id, setup_version_id=setup_version_id, config=config)
66
+ channel = self._init_channel(client_config)
67
+ self.stub = cost_service_pb2_grpc.CostServiceStub(channel)
68
+ logger.debug("Channel client 'Cost' initialized succesfully")
69
+
70
+ def add(
71
+ self,
72
+ name: str,
73
+ cost_config_name: str,
74
+ quantity: float,
75
+ ) -> None:
76
+ """Create a new record in the cost database.
77
+
78
+ Args:
79
+ name: The name of the cost
80
+ cost_config_name: The name of the cost config
81
+ quantity: The quantity of the cost
65
82
 
66
- request = cost_pb2.AddCostRequest(**valid_data.model_dump())
67
- return self.exec_grpc_query("AddCost", request)
83
+ Raises:
84
+ CostServiceError: If the cost config is invalid
85
+ """
86
+ with self._handle_grpc_errors("AddCost"):
87
+ cost_config = self.config.get(cost_config_name)
88
+ if cost_config is None:
89
+ msg = f"Cost config {cost_config_name} not found in the configuration."
90
+ logger.error(msg)
91
+ raise CostServiceError(msg)
92
+ valid_data = CostData.model_validate({
93
+ "name": name,
94
+ "cost": cost_config.rate * quantity,
95
+ "unit": cost_config.unit,
96
+ "cost_type": CostType[cost_config.cost_type],
97
+ "mission_id": self.mission_id,
98
+ "rate": cost_config.rate,
99
+ "quantity": quantity,
100
+ "setup_version_id": self.setup_version_id,
101
+ })
102
+ request = cost_pb2.AddCostRequest(
103
+ cost=valid_data.cost,
104
+ name=valid_data.name,
105
+ unit=valid_data.unit,
106
+ cost_type=valid_data.cost_type.name,
107
+ mission_id=valid_data.mission_id,
108
+ rate=valid_data.rate,
109
+ quantity=valid_data.quantity,
110
+ setup_version_id=valid_data.setup_version_id,
111
+ )
112
+ self.exec_grpc_query("AddCost", request)
113
+ logger.debug("Cost added with cost_dict: %s", valid_data.model_dump())
114
+
115
+ def get(self, name: str) -> list[CostData]:
116
+ """Get a record from the database.
117
+
118
+ Args:
119
+ name: The name of the cost
68
120
 
69
- def get(self, cost_dict: dict[str, Any]) -> list[CostData]:
70
- """Get records from the database.
121
+ Returns:
122
+ CostData: The cost data
123
+ """
124
+ with self._handle_grpc_errors("GetCost"):
125
+ request = cost_pb2.GetCostRequest(name=name, mission_id=self.mission_id)
126
+ response: cost_pb2.GetCostResponse = self.exec_grpc_query("GetCost", request)
127
+ cost_data_list = [
128
+ json_format.MessageToDict(
129
+ cost,
130
+ preserving_proto_field_name=True,
131
+ always_print_fields_with_no_presence=True,
132
+ )
133
+ for cost in response.costs
134
+ ]
135
+ logger.debug("Costs retrieved with cost_dict: %s", cost_data_list)
136
+ return [CostData.model_validate(cost_data) for cost_data in cost_data_list]
137
+
138
+ def get_filtered(
139
+ self,
140
+ names: list[str] | None = None,
141
+ cost_types: list[Literal["TOKEN_INPUT", "TOKEN_OUTPUT", "API_CALL", "STORAGE", "TIME", "OTHER"]] | None = None,
142
+ ) -> list[CostData]:
143
+ """Get a list of records from the database.
144
+
145
+ Args:
146
+ names: The names of the costs
147
+ cost_types: The types of the costs
71
148
 
72
149
  Returns:
73
- list[CostData]: The list of records
150
+ list[CostData]: The cost data
74
151
  """
75
- if "mission_id" not in cost_dict:
76
- return []
77
- if "name" in cost_dict:
78
- return self._get_costs_by_name(cost_dict)
79
- if "type" in cost_dict:
80
- return self._get_costs_by_type(cost_dict)
81
- return self._get_costs_by_mission(cost_dict)
152
+ with self._handle_grpc_errors("GetCosts"):
153
+ request = cost_pb2.GetCostsRequest(
154
+ mission_id=self.mission_id,
155
+ filter=cost_pb2.CostFilter(
156
+ names=names or [],
157
+ cost_types=cost_types or [],
158
+ ),
159
+ )
160
+ response: cost_pb2.GetCostsResponse = self.exec_grpc_query("GetCosts", request)
161
+ cost_data_list = [
162
+ json_format.MessageToDict(
163
+ cost,
164
+ preserving_proto_field_name=True,
165
+ always_print_fields_with_no_presence=True,
166
+ )
167
+ for cost in response.costs
168
+ ]
169
+ logger.debug("Filtered costs retrieved with cost_dict: %s", cost_data_list)
170
+ return [CostData.model_validate(cost_data) for cost_data in cost_data_list]
@@ -1,10 +1,10 @@
1
1
  """Default filesystem."""
2
2
 
3
- import logging
4
3
  import os
5
4
  import tempfile
6
5
  from pathlib import Path
7
6
 
7
+ from digitalkin.logger import logger
8
8
  from digitalkin.services.filesystem.filesystem_strategy import (
9
9
  FilesystemData,
10
10
  FilesystemServiceError,
@@ -12,20 +12,19 @@ from digitalkin.services.filesystem.filesystem_strategy import (
12
12
  FileType,
13
13
  )
14
14
 
15
- logger = logging.getLogger(__name__)
16
-
17
15
 
18
16
  class DefaultFilesystem(FilesystemStrategy):
19
17
  """Default state filesystem strategy."""
20
18
 
21
- def __init__(self, mission_id: str, config: dict[str, str]) -> None:
19
+ def __init__(self, mission_id: str, setup_version_id: str, config: dict[str, str]) -> None:
22
20
  """Initialize the default filesystem strategy.
23
21
 
24
22
  Args:
25
23
  mission_id: The ID of the mission this strategy is associated with
24
+ setup_version_id: The ID of the setup version this strategy is associated with
26
25
  config: A dictionary mapping names to Pydantic model classes
27
26
  """
28
- super().__init__(mission_id, config)
27
+ super().__init__(mission_id, setup_version_id, config)
29
28
  self.temp_root: str = self.config.get("temp_root", "") or tempfile.gettempdir()
30
29
  os.makedirs(self.temp_root, exist_ok=True)
31
30
  self.db: dict[str, FilesystemData] = {}
@@ -176,7 +175,7 @@ class DefaultFilesystem(FilesystemStrategy):
176
175
  try:
177
176
  os.remove(file_path)
178
177
  del self.db[name]
179
- logger.info("File %s successfully deleted.", name)
178
+ logger.debug("File %s successfully deleted.", name)
180
179
 
181
180
  except OSError:
182
181
  msg = f"Error deleting file {name} from filesystem"
@@ -35,14 +35,15 @@ class FilesystemData(BaseModel):
35
35
  class FilesystemStrategy(BaseStrategy, ABC):
36
36
  """Abstract base class for filesystem strategies."""
37
37
 
38
- def __init__(self, mission_id: str, config: dict[str, str]) -> None:
38
+ def __init__(self, mission_id: str, setup_version_id: str, config: dict[str, str]) -> None:
39
39
  """Initialize the strategy.
40
40
 
41
41
  Args:
42
42
  mission_id: The ID of the mission this strategy is associated with
43
+ setup_version_id: The ID of the setup version this strategy is associated with
43
44
  config: configuration dictionary for the filesystem strategy
44
45
  """
45
- super().__init__(mission_id)
46
+ super().__init__(mission_id, setup_version_id)
46
47
  self.config: dict[str, str] = config
47
48
 
48
49
  @abstractmethod
@@ -1,16 +1,21 @@
1
1
  """Grpc filesystem."""
2
2
 
3
- import logging
4
3
  from collections.abc import Generator
5
4
  from contextlib import contextmanager
6
5
  from typing import Any
7
6
 
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
7
+ from digitalkin_proto.digitalkin.filesystem.v2 import (
8
+ filesystem_pb2,
9
+ filesystem_service_pb2_grpc,
10
+ )
11
+ from digitalkin_proto.digitalkin.filesystem.v2.filesystem_pb2 import (
12
+ FileType as FileTypeProto,
13
+ )
10
14
 
11
15
  from digitalkin.grpc_servers.utils.exceptions import ServerError
12
16
  from digitalkin.grpc_servers.utils.grpc_client_wrapper import GrpcClientWrapper
13
17
  from digitalkin.grpc_servers.utils.models import ClientConfig
18
+ from digitalkin.logger import logger
14
19
  from digitalkin.services.filesystem.filesystem_strategy import (
15
20
  FilesystemData,
16
21
  FilesystemServiceError,
@@ -18,32 +23,10 @@ from digitalkin.services.filesystem.filesystem_strategy import (
18
23
  FileType,
19
24
  )
20
25
 
21
- logger = logging.getLogger(__name__)
22
-
23
26
 
24
27
  class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper):
25
28
  """Default state filesystem strategy."""
26
29
 
27
- def __init__(
28
- self,
29
- mission_id: str,
30
- config: dict[str, str],
31
- client_config: ClientConfig,
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
- client_config: The client 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(client_config)
44
- self.stub = filesystem_service_pb2_grpc.FilesystemServiceStub(channel)
45
- logger.info("Channel client 'Filesystem' initialized succesfully")
46
-
47
30
  @staticmethod
48
31
  @contextmanager
49
32
  def _handle_grpc_errors(operation: str) -> Generator[Any, Any, Any]:
@@ -56,7 +39,7 @@ class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper):
56
39
  operation: Description of the operation being performed.
57
40
 
58
41
  Raises:
59
- ValueError: Error wiht the model validation.
42
+ ValueError: Error with the model validation.
60
43
  ServerError: from gRPC Client.
61
44
  FilesystemServiceError: Filesystem service internal.
62
45
  """
@@ -71,6 +54,28 @@ class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper):
71
54
  logger.exception(msg)
72
55
  raise FilesystemServiceError(msg) from e
73
56
 
57
+ def __init__(
58
+ self,
59
+ mission_id: str,
60
+ setup_version_id: str,
61
+ config: dict[str, str],
62
+ client_config: ClientConfig,
63
+ **kwargs, # noqa: ANN003, ARG002
64
+ ) -> None:
65
+ """Initialize the default filesystem strategy.
66
+
67
+ Args:
68
+ mission_id: The ID of the mission this strategy is associated with
69
+ setup_version_id: The ID of the setup version this strategy is associated with
70
+ config: A dictionary mapping names to Pydantic model classes
71
+ client_config: The server configuration object
72
+ kwargs: other optional arguments to pass to the parent class constructor
73
+ """
74
+ super().__init__(mission_id, setup_version_id, config)
75
+ channel = self._init_channel(client_config)
76
+ self.stub = filesystem_service_pb2_grpc.FilesystemServiceStub(channel)
77
+ logger.info("Channel client 'Filesystem' initialized succesfully")
78
+
74
79
  def upload(self, content: bytes, name: str, file_type: FileType) -> FilesystemData:
75
80
  """Create a new file in the file system.
76
81
 
@@ -5,7 +5,7 @@ from typing import Any, ClassVar
5
5
  from pydantic import BaseModel, Field, PrivateAttr
6
6
 
7
7
  from digitalkin.services.agent import AgentStrategy, DefaultAgent
8
- from digitalkin.services.cost import CostStrategy, DefaultCost
8
+ from digitalkin.services.cost import CostStrategy, DefaultCost, GrpcCost
9
9
  from digitalkin.services.filesystem import DefaultFilesystem, FilesystemStrategy, GrpcFilesystem
10
10
  from digitalkin.services.identity import DefaultIdentity, IdentityStrategy
11
11
  from digitalkin.services.registry import DefaultRegistry, RegistryStrategy
@@ -30,7 +30,7 @@ class ServicesConfig(BaseModel):
30
30
  )
31
31
  _config_storage: dict[str, Any | None] = PrivateAttr(default_factory=dict)
32
32
  _cost: ServicesStrategy[CostStrategy] = PrivateAttr(
33
- default_factory=lambda: ServicesStrategy(local=DefaultCost, remote=DefaultCost)
33
+ default_factory=lambda: ServicesStrategy(local=DefaultCost, remote=GrpcCost)
34
34
  )
35
35
  _config_cost: dict[str, Any | None] = PrivateAttr(default_factory=dict)
36
36
  _snapshot: ServicesStrategy[SnapshotStrategy] = PrivateAttr(
@@ -111,12 +111,13 @@ class ServicesConfig(BaseModel):
111
111
  """
112
112
  return getattr(self, f"_config_{name}", {})
113
113
 
114
- def init_strategy(self, name: str, mission_id: str) -> ServicesStrategy:
114
+ def init_strategy(self, name: str, mission_id: str, setup_version_id: str) -> ServicesStrategy:
115
115
  """Initialize a specific strategy.
116
116
 
117
117
  Args:
118
118
  name: The name of the strategy to initialize
119
119
  mission_id: The ID of the mission this strategy is associated with
120
+ setup_version_id: The setup version ID for the strategy
120
121
 
121
122
  Returns:
122
123
  The initialized strategy instance
@@ -129,8 +130,8 @@ class ServicesConfig(BaseModel):
129
130
  msg = f"Strategy {name} not found in ServicesConfig."
130
131
  raise ValueError(msg)
131
132
 
132
- # Instantiate the strategy with the mission ID and configuration
133
- return strategy_type(mission_id, **self.get_strategy_config(name) or {})
133
+ # Instantiate the strategy with the mission ID, setup version ID, and configuration
134
+ return strategy_type(mission_id, setup_version_id, **self.get_strategy_config(name) or {})
134
135
 
135
136
  @property
136
137
  def storage(self) -> type[StorageStrategy]:
@@ -0,0 +1 @@
1
+ """This module is responsible for handling the setup service."""