digitalkin 0.2.10__tar.gz → 0.2.12__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {digitalkin-0.2.10 → digitalkin-0.2.12}/PKG-INFO +2 -2
- {digitalkin-0.2.10 → digitalkin-0.2.12}/examples/modules/storage_module.py +2 -2
- {digitalkin-0.2.10 → digitalkin-0.2.12}/pyproject.toml +6 -3
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/__version__.py +1 -1
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/grpc_servers/module_servicer.py +1 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/modules/_base_module.py +3 -1
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/modules/job_manager.py +2 -1
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/base_strategy.py +3 -1
- digitalkin-0.2.12/src/digitalkin/services/cost/cost_strategy.py +93 -0
- digitalkin-0.2.12/src/digitalkin/services/cost/default_cost.py +109 -0
- digitalkin-0.2.12/src/digitalkin/services/cost/grpc_cost.py +162 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/filesystem/default_filesystem.py +3 -2
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/filesystem/filesystem_strategy.py +3 -2
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/filesystem/grpc_filesystem.py +23 -21
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/services_config.py +6 -5
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/storage/default_storage.py +2 -1
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/storage/grpc_storage.py +3 -2
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/storage/storage_strategy.py +10 -9
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin.egg-info/PKG-INFO +2 -2
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin.egg-info/requires.txt +1 -1
- digitalkin-0.2.10/src/digitalkin/services/cost/cost_strategy.py +0 -45
- digitalkin-0.2.10/src/digitalkin/services/cost/default_cost.py +0 -30
- digitalkin-0.2.10/src/digitalkin/services/cost/grpc_cost.py +0 -81
- {digitalkin-0.2.10 → digitalkin-0.2.12}/LICENSE +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/README.md +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/examples/base_server/__init__.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/examples/base_server/mock/__init__.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/examples/base_server/mock/mock_pb2.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/examples/base_server/mock/mock_pb2_grpc.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/examples/base_server/server_async_insecure.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/examples/base_server/server_async_secure.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/examples/base_server/server_sync_insecure.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/examples/base_server/server_sync_secure.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/examples/modules/__init__.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/examples/modules/minimal_llm_module.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/examples/modules/text_transform_module.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/setup.cfg +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/__init__.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/grpc_servers/__init__.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/grpc_servers/_base_server.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/grpc_servers/module_server.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/grpc_servers/registry_server.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/grpc_servers/registry_servicer.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/grpc_servers/utils/exceptions.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/grpc_servers/utils/factory.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/grpc_servers/utils/grpc_client_wrapper.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/grpc_servers/utils/models.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/grpc_servers/utils/types.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/logger.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/models/__init__.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/models/module/__init__.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/models/module/module.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/models/module/module_types.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/models/services/__init__.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/models/services/cost.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/models/services/storage.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/modules/__init__.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/modules/archetype_module.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/modules/tool_module.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/modules/trigger_module.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/py.typed +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/__init__.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/agent/__init__.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/agent/agent_strategy.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/agent/default_agent.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/cost/__init__.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/filesystem/__init__.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/identity/__init__.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/identity/default_identity.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/identity/identity_strategy.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/registry/__init__.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/registry/default_registry.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/registry/registry_strategy.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/services_models.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/setup/default_setup.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/setup/grpc_setup.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/setup/setup_strategy.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/snapshot/__init__.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/snapshot/default_snapshot.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/snapshot/snapshot_strategy.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/storage/__init__.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/utils/__init__.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/utils/arg_parser.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/utils/llm_ready_schema.py +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin.egg-info/SOURCES.txt +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin.egg-info/dependency_links.txt +0 -0
- {digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin.egg-info/top_level.txt +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: digitalkin
|
|
3
|
-
Version: 0.2.
|
|
3
|
+
Version: 0.2.12
|
|
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,7 +452,7 @@ Classifier: License :: Other/Proprietary License
|
|
|
452
452
|
Requires-Python: >=3.10
|
|
453
453
|
Description-Content-Type: text/markdown
|
|
454
454
|
License-File: LICENSE
|
|
455
|
-
Requires-Dist: digitalkin-proto>=0.1.
|
|
455
|
+
Requires-Dist: digitalkin-proto>=0.1.10
|
|
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
|
|
@@ -62,7 +62,7 @@ class ExampleModule(ArchetypeModule[ExampleInput, ExampleOutput, ExampleSetup, E
|
|
|
62
62
|
|
|
63
63
|
# Define services_config_params with default values
|
|
64
64
|
services_config_strategies = {}
|
|
65
|
-
services_config_params = {"storage": {"config": {"
|
|
65
|
+
services_config_params = {"storage": {"config": {"example": ExampleOutput}}, "filesystem": {"config": {}}}
|
|
66
66
|
|
|
67
67
|
def __init__(self, job_id: str, mission_id: str) -> None:
|
|
68
68
|
"""Initialize the example module.
|
|
@@ -172,7 +172,7 @@ async def test_module() -> None:
|
|
|
172
172
|
def test_storage_directly() -> None:
|
|
173
173
|
"""Test the storage service directly."""
|
|
174
174
|
# Initialize storage service
|
|
175
|
-
storage = ServicesConfig().storage(mission_id="test-mission", config={"
|
|
175
|
+
storage = ServicesConfig().storage(mission_id="test-mission", config={"example": ExampleStorage})
|
|
176
176
|
|
|
177
177
|
# Create a test record
|
|
178
178
|
storage.store("example", "test_table", {"test_key": "test_value"}, "OUTPUT")
|
|
@@ -12,7 +12,7 @@
|
|
|
12
12
|
|
|
13
13
|
keywords = [ "digitalkin", "kin", "agent", "gprc", "sdk" ]
|
|
14
14
|
# Version of the package automatically updated by bump2version (that is why it is separated)
|
|
15
|
-
version = "0.2.
|
|
15
|
+
version = "0.2.12"
|
|
16
16
|
|
|
17
17
|
classifiers = [
|
|
18
18
|
"Development Status :: 3 - Alpha",
|
|
@@ -29,13 +29,13 @@
|
|
|
29
29
|
]
|
|
30
30
|
|
|
31
31
|
dependencies = [
|
|
32
|
-
"digitalkin-proto>=0.1.
|
|
32
|
+
"digitalkin-proto>=0.1.10",
|
|
33
33
|
"grpcio-health-checking>=1.71.0",
|
|
34
34
|
"grpcio-reflection>=1.71.0",
|
|
35
35
|
"grpcio-status>=1.71.0",
|
|
36
36
|
"openai>=1.76.2",
|
|
37
37
|
"pydantic>=2.11.4",
|
|
38
|
-
|
|
38
|
+
]
|
|
39
39
|
|
|
40
40
|
[project.optional-dependencies]
|
|
41
41
|
dev = [
|
|
@@ -182,6 +182,9 @@
|
|
|
182
182
|
"ARG002",
|
|
183
183
|
]
|
|
184
184
|
|
|
185
|
+
[tool.ruff.lint.pylint]
|
|
186
|
+
max-args = 10
|
|
187
|
+
|
|
185
188
|
[tool.ruff.format]
|
|
186
189
|
quote-style = "double"
|
|
187
190
|
indent-style = "space"
|
|
@@ -96,6 +96,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer):
|
|
|
96
96
|
input_data,
|
|
97
97
|
setup_data,
|
|
98
98
|
mission_id=request.mission_id,
|
|
99
|
+
setup_version_id=setup_data_class.current_setup_version.id,
|
|
99
100
|
callback=self.add_to_queue,
|
|
100
101
|
)
|
|
101
102
|
job_id, module = result
|
|
@@ -48,17 +48,19 @@ class BaseModule(ABC, Generic[InputModelT, OutputModelT, SetupModelT, SecretMode
|
|
|
48
48
|
def _init_strategies(self) -> None:
|
|
49
49
|
"""Initialize the services configuration."""
|
|
50
50
|
for service_name in self.services_config.valid_strategy_names():
|
|
51
|
-
service = self.services_config.init_strategy(service_name, self.mission_id)
|
|
51
|
+
service = self.services_config.init_strategy(service_name, self.mission_id, self.setup_version_id)
|
|
52
52
|
setattr(self, service_name, service)
|
|
53
53
|
|
|
54
54
|
def __init__(
|
|
55
55
|
self,
|
|
56
56
|
job_id: str,
|
|
57
57
|
mission_id: str,
|
|
58
|
+
setup_version_id: str,
|
|
58
59
|
) -> None:
|
|
59
60
|
"""Initialize the module."""
|
|
60
61
|
self.job_id: str = job_id
|
|
61
62
|
self.mission_id: str = mission_id
|
|
63
|
+
self.setup_version_id: str = setup_version_id
|
|
62
64
|
self._status = ModuleStatus.CREATED
|
|
63
65
|
self._task: asyncio.Task | None = None
|
|
64
66
|
# Initialize services configuration
|
|
@@ -79,6 +79,7 @@ class JobManager(ArgParser):
|
|
|
79
79
|
input_data: InputModelT,
|
|
80
80
|
setup_data: SetupModelT,
|
|
81
81
|
mission_id: str,
|
|
82
|
+
setup_version_id: str,
|
|
82
83
|
callback: Callable[[str, OutputModelT], Coroutine[Any, Any, None]],
|
|
83
84
|
) -> tuple[str, BaseModule[InputModelT, OutputModelT, SetupModelT, SecretModelT]]: # type: ignore
|
|
84
85
|
"""Start new module job in background (asyncio).
|
|
@@ -94,7 +95,7 @@ class JobManager(ArgParser):
|
|
|
94
95
|
job_id = str(uuid.uuid4())
|
|
95
96
|
"""TODO: check uniqueness of the job_id"""
|
|
96
97
|
# Création et démarrage du module
|
|
97
|
-
module = self.module_class(job_id, mission_id=mission_id)
|
|
98
|
+
module = self.module_class(job_id, mission_id=mission_id, setup_version_id=setup_version_id)
|
|
98
99
|
self.modules[job_id] = module
|
|
99
100
|
try:
|
|
100
101
|
await module.start(input_data, setup_data, await JobManager._job_specific_callback(callback, job_id))
|
|
@@ -9,10 +9,12 @@ class BaseStrategy(ABC):
|
|
|
9
9
|
This class defines the interface for all strategies.
|
|
10
10
|
"""
|
|
11
11
|
|
|
12
|
-
def __init__(self, mission_id: str) -> None:
|
|
12
|
+
def __init__(self, mission_id: str, setup_version_id: str) -> None:
|
|
13
13
|
"""Initialize the strategy.
|
|
14
14
|
|
|
15
15
|
Args:
|
|
16
16
|
mission_id: The ID of the mission this strategy is associated with
|
|
17
|
+
setup_version_id: The ID of the setup version this strategy is associated with
|
|
17
18
|
"""
|
|
18
19
|
self.mission_id: str = mission_id
|
|
20
|
+
self.setup_version_id: str = setup_version_id
|
|
@@ -0,0 +1,93 @@
|
|
|
1
|
+
"""This module contains the abstract base class for cost strategies."""
|
|
2
|
+
|
|
3
|
+
from abc import ABC, abstractmethod
|
|
4
|
+
from enum import Enum
|
|
5
|
+
from typing import Literal
|
|
6
|
+
|
|
7
|
+
from pydantic import BaseModel
|
|
8
|
+
|
|
9
|
+
from digitalkin.services.base_strategy import BaseStrategy
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
class CostType(Enum):
|
|
13
|
+
"""Enum defining the types of costs that can be registered."""
|
|
14
|
+
|
|
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
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
class CostData(BaseModel):
|
|
41
|
+
"""Data model for cost operations."""
|
|
42
|
+
|
|
43
|
+
cost: float
|
|
44
|
+
mission_id: str
|
|
45
|
+
name: str
|
|
46
|
+
cost_type: CostType
|
|
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."""
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
class CostStrategy(BaseStrategy, ABC):
|
|
58
|
+
"""Abstract base class for cost strategies."""
|
|
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
|
+
|
|
71
|
+
@abstractmethod
|
|
72
|
+
def add(
|
|
73
|
+
self,
|
|
74
|
+
name: str,
|
|
75
|
+
cost_config_name: str,
|
|
76
|
+
quantity: float,
|
|
77
|
+
) -> None:
|
|
78
|
+
"""Register a new cost."""
|
|
79
|
+
|
|
80
|
+
@abstractmethod
|
|
81
|
+
def get(
|
|
82
|
+
self,
|
|
83
|
+
name: str,
|
|
84
|
+
) -> list[CostData]:
|
|
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."""
|
|
@@ -0,0 +1,109 @@
|
|
|
1
|
+
"""Default cost."""
|
|
2
|
+
|
|
3
|
+
import logging
|
|
4
|
+
from typing import Literal
|
|
5
|
+
|
|
6
|
+
from digitalkin.services.cost.cost_strategy import CostConfig, CostData, CostServiceError, CostStrategy, CostType
|
|
7
|
+
|
|
8
|
+
logger = logging.getLogger(__name__)
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
class DefaultCost(CostStrategy):
|
|
12
|
+
"""Default cost strategy."""
|
|
13
|
+
|
|
14
|
+
def __init__(self, mission_id: str, setup_version_id: str, config: dict[str, CostConfig]) -> None:
|
|
15
|
+
"""Initialize the strategy.
|
|
16
|
+
|
|
17
|
+
Args:
|
|
18
|
+
mission_id: The ID of the mission this strategy is associated with
|
|
19
|
+
setup_version_id: The ID of the setup version this strategy is associated with
|
|
20
|
+
config: The configuration dictionary for the cost
|
|
21
|
+
"""
|
|
22
|
+
super().__init__(mission_id=mission_id, setup_version_id=setup_version_id, config=config)
|
|
23
|
+
self.db: dict[str, list[CostData]] = {}
|
|
24
|
+
|
|
25
|
+
def add(
|
|
26
|
+
self,
|
|
27
|
+
name: str,
|
|
28
|
+
cost_config_name: str,
|
|
29
|
+
quantity: float,
|
|
30
|
+
) -> None:
|
|
31
|
+
"""Create a new record in the cost database.
|
|
32
|
+
|
|
33
|
+
Args:
|
|
34
|
+
name: The name of the cost
|
|
35
|
+
cost_config_name: The name of the cost config
|
|
36
|
+
quantity: The quantity of the cost
|
|
37
|
+
|
|
38
|
+
Raises:
|
|
39
|
+
CostServiceError: If the cost data is invalid or if the cost already exists
|
|
40
|
+
"""
|
|
41
|
+
cost_config = self.config.get(cost_config_name)
|
|
42
|
+
if cost_config is None:
|
|
43
|
+
msg = f"Cost config {cost_config_name} not found in the configuration."
|
|
44
|
+
logger.error(msg)
|
|
45
|
+
raise CostServiceError(msg)
|
|
46
|
+
cost_data = CostData.model_validate({
|
|
47
|
+
"name": name,
|
|
48
|
+
"cost": cost_config.rate * quantity,
|
|
49
|
+
"unit": cost_config.unit,
|
|
50
|
+
"cost_type": getattr(CostType, cost_config.cost_type),
|
|
51
|
+
"mission_id": self.mission_id,
|
|
52
|
+
"rate": cost_config.rate,
|
|
53
|
+
"quantity": quantity,
|
|
54
|
+
"setup_version_id": self.setup_version_id,
|
|
55
|
+
})
|
|
56
|
+
if cost_data.mission_id not in self.db:
|
|
57
|
+
self.db[cost_data.mission_id] = []
|
|
58
|
+
if cost_data.name in [cost.name for cost in self.db[cost_data.mission_id]]:
|
|
59
|
+
msg = f"Cost with name {cost_data.name} already exists in mission {cost_data.mission_id}"
|
|
60
|
+
logger.error(msg)
|
|
61
|
+
raise CostServiceError(msg)
|
|
62
|
+
self.db[cost_data.mission_id].append(cost_data)
|
|
63
|
+
|
|
64
|
+
def get(self, name: str) -> list[CostData]:
|
|
65
|
+
"""Get a record from the database.
|
|
66
|
+
|
|
67
|
+
Args:
|
|
68
|
+
name: The name of the cost
|
|
69
|
+
|
|
70
|
+
Returns:
|
|
71
|
+
list[CostData]: The cost data
|
|
72
|
+
|
|
73
|
+
Raises:
|
|
74
|
+
CostServiceError: If the cost data is invalid or if the cost does not exist
|
|
75
|
+
"""
|
|
76
|
+
if self.mission_id not in self.db:
|
|
77
|
+
msg = f"Mission {self.mission_id} not found in the database."
|
|
78
|
+
logger.warning(msg)
|
|
79
|
+
raise CostServiceError(msg)
|
|
80
|
+
|
|
81
|
+
return [cost for cost in self.db[self.mission_id] if cost.name == name] or []
|
|
82
|
+
|
|
83
|
+
def get_filtered(
|
|
84
|
+
self,
|
|
85
|
+
names: list[str] | None = None,
|
|
86
|
+
cost_types: list[Literal["TOKEN_INPUT", "TOKEN_OUTPUT", "API_CALL", "STORAGE", "TIME", "OTHER"]] | None = None,
|
|
87
|
+
) -> list[CostData]:
|
|
88
|
+
"""Get records from the database.
|
|
89
|
+
|
|
90
|
+
Args:
|
|
91
|
+
names: The names of the costs
|
|
92
|
+
cost_types: The types of the costs
|
|
93
|
+
|
|
94
|
+
Returns:
|
|
95
|
+
list[CostData]: The list of records
|
|
96
|
+
|
|
97
|
+
Raises:
|
|
98
|
+
CostServiceError: If the cost data is invalid or if the cost does not exist
|
|
99
|
+
"""
|
|
100
|
+
if self.mission_id not in self.db:
|
|
101
|
+
msg = f"Mission {self.mission_id} not found in the database."
|
|
102
|
+
logger.warning(msg)
|
|
103
|
+
raise CostServiceError(msg)
|
|
104
|
+
|
|
105
|
+
return [
|
|
106
|
+
cost
|
|
107
|
+
for cost in self.db[self.mission_id]
|
|
108
|
+
if (names and cost.name in names) or (cost_types and cost.cost_type in cost_types)
|
|
109
|
+
]
|
|
@@ -0,0 +1,162 @@
|
|
|
1
|
+
"""This module implements the default Cost strategy."""
|
|
2
|
+
|
|
3
|
+
import logging
|
|
4
|
+
from collections.abc import Generator
|
|
5
|
+
from contextlib import contextmanager
|
|
6
|
+
from typing import Any, Literal
|
|
7
|
+
|
|
8
|
+
from digitalkin_proto.digitalkin.cost.v1 import cost_pb2, cost_service_pb2_grpc
|
|
9
|
+
from google.protobuf import json_format
|
|
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 ClientConfig
|
|
14
|
+
from digitalkin.services.cost.cost_strategy import CostConfig, CostData, CostServiceError, CostStrategy, CostType
|
|
15
|
+
|
|
16
|
+
logger = logging.getLogger(__name__)
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
class GrpcCost(CostStrategy, GrpcClientWrapper):
|
|
20
|
+
"""This class implements the default Cost strategy."""
|
|
21
|
+
|
|
22
|
+
@staticmethod
|
|
23
|
+
@contextmanager
|
|
24
|
+
def _handle_grpc_errors(operation: str) -> Generator[Any, Any, Any]:
|
|
25
|
+
"""Context manager for consistent gRPC error handling.
|
|
26
|
+
|
|
27
|
+
Yields:
|
|
28
|
+
Allow error handling in context.
|
|
29
|
+
|
|
30
|
+
Args:
|
|
31
|
+
operation: Description of the operation being performed.
|
|
32
|
+
|
|
33
|
+
Raises:
|
|
34
|
+
ValueError: Error with the model validation.
|
|
35
|
+
ServerError: from gRPC Client.
|
|
36
|
+
CostServiceError: Unexpected error.
|
|
37
|
+
"""
|
|
38
|
+
try:
|
|
39
|
+
yield
|
|
40
|
+
except CostServiceError as e:
|
|
41
|
+
msg = f"CostServiceError in {operation}: {e}"
|
|
42
|
+
logger.exception(msg)
|
|
43
|
+
raise CostServiceError(msg) from e
|
|
44
|
+
except ServerError as e:
|
|
45
|
+
msg = f"gRPC {operation} failed: {e}"
|
|
46
|
+
logger.exception(msg)
|
|
47
|
+
raise ServerError(msg) from e
|
|
48
|
+
except Exception as e:
|
|
49
|
+
msg = f"Unexpected error in {operation}"
|
|
50
|
+
logger.exception(msg)
|
|
51
|
+
raise CostServiceError(msg) from e
|
|
52
|
+
|
|
53
|
+
def __init__(
|
|
54
|
+
self, mission_id: str, setup_version_id: str, config: dict[str, CostConfig], client_config: ClientConfig
|
|
55
|
+
) -> None:
|
|
56
|
+
"""Initialize the cost."""
|
|
57
|
+
super().__init__(mission_id=mission_id, setup_version_id=setup_version_id, config=config)
|
|
58
|
+
channel = self._init_channel(client_config)
|
|
59
|
+
self.stub = cost_service_pb2_grpc.CostServiceStub(channel)
|
|
60
|
+
logger.info("Channel client 'Cost' initialized succesfully")
|
|
61
|
+
|
|
62
|
+
def add(
|
|
63
|
+
self,
|
|
64
|
+
name: str,
|
|
65
|
+
cost_config_name: str,
|
|
66
|
+
quantity: float,
|
|
67
|
+
) -> None:
|
|
68
|
+
"""Create a new record in the cost database.
|
|
69
|
+
|
|
70
|
+
Args:
|
|
71
|
+
name: The name of the cost
|
|
72
|
+
cost_config_name: The name of the cost config
|
|
73
|
+
quantity: The quantity of the cost
|
|
74
|
+
|
|
75
|
+
Raises:
|
|
76
|
+
CostServiceError: If the cost config is invalid
|
|
77
|
+
"""
|
|
78
|
+
with self._handle_grpc_errors("AddCost"):
|
|
79
|
+
cost_config = self.config.get(cost_config_name)
|
|
80
|
+
if cost_config is None:
|
|
81
|
+
msg = f"Cost config {cost_config_name} not found in the configuration."
|
|
82
|
+
logger.error(msg)
|
|
83
|
+
raise CostServiceError(msg)
|
|
84
|
+
valid_data = CostData.model_validate({
|
|
85
|
+
"name": name,
|
|
86
|
+
"cost": cost_config.rate * quantity,
|
|
87
|
+
"unit": cost_config.unit,
|
|
88
|
+
"cost_type": CostType[cost_config.cost_type],
|
|
89
|
+
"mission_id": self.mission_id,
|
|
90
|
+
"rate": cost_config.rate,
|
|
91
|
+
"quantity": quantity,
|
|
92
|
+
"setup_version_id": self.setup_version_id,
|
|
93
|
+
})
|
|
94
|
+
request = cost_pb2.AddCostRequest(
|
|
95
|
+
cost=valid_data.cost,
|
|
96
|
+
name=valid_data.name,
|
|
97
|
+
unit=valid_data.unit,
|
|
98
|
+
cost_type=valid_data.cost_type.name,
|
|
99
|
+
mission_id=valid_data.mission_id,
|
|
100
|
+
rate=valid_data.rate,
|
|
101
|
+
quantity=valid_data.quantity,
|
|
102
|
+
setup_version_id=valid_data.setup_version_id,
|
|
103
|
+
)
|
|
104
|
+
self.exec_grpc_query("AddCost", request)
|
|
105
|
+
logger.debug("Cost added with cost_dict: %s", valid_data.model_dump())
|
|
106
|
+
|
|
107
|
+
def get(self, name: str) -> list[CostData]:
|
|
108
|
+
"""Get a record from the database.
|
|
109
|
+
|
|
110
|
+
Args:
|
|
111
|
+
name: The name of the cost
|
|
112
|
+
|
|
113
|
+
Returns:
|
|
114
|
+
CostData: The cost data
|
|
115
|
+
"""
|
|
116
|
+
with self._handle_grpc_errors("GetCost"):
|
|
117
|
+
request = cost_pb2.GetCostRequest(name=name, mission_id=self.mission_id)
|
|
118
|
+
response: cost_pb2.GetCostResponse = self.exec_grpc_query("GetCost", request)
|
|
119
|
+
cost_data_list = [
|
|
120
|
+
json_format.MessageToDict(
|
|
121
|
+
cost,
|
|
122
|
+
preserving_proto_field_name=True,
|
|
123
|
+
always_print_fields_with_no_presence=True,
|
|
124
|
+
)
|
|
125
|
+
for cost in response.costs
|
|
126
|
+
]
|
|
127
|
+
logger.debug("Costs retrieved with cost_dict: %s", cost_data_list)
|
|
128
|
+
return [CostData.model_validate(cost_data) for cost_data in cost_data_list]
|
|
129
|
+
|
|
130
|
+
def get_filtered(
|
|
131
|
+
self,
|
|
132
|
+
names: list[str] | None = None,
|
|
133
|
+
cost_types: list[Literal["TOKEN_INPUT", "TOKEN_OUTPUT", "API_CALL", "STORAGE", "TIME", "OTHER"]] | None = None,
|
|
134
|
+
) -> list[CostData]:
|
|
135
|
+
"""Get a list of records from the database.
|
|
136
|
+
|
|
137
|
+
Args:
|
|
138
|
+
names: The names of the costs
|
|
139
|
+
cost_types: The types of the costs
|
|
140
|
+
|
|
141
|
+
Returns:
|
|
142
|
+
list[CostData]: The cost data
|
|
143
|
+
"""
|
|
144
|
+
with self._handle_grpc_errors("GetCosts"):
|
|
145
|
+
request = cost_pb2.GetCostsRequest(
|
|
146
|
+
mission_id=self.mission_id,
|
|
147
|
+
filter=cost_pb2.CostFilter(
|
|
148
|
+
names=names or [],
|
|
149
|
+
cost_types=cost_types or [],
|
|
150
|
+
),
|
|
151
|
+
)
|
|
152
|
+
response: cost_pb2.GetCostsResponse = self.exec_grpc_query("GetCosts", request)
|
|
153
|
+
cost_data_list = [
|
|
154
|
+
json_format.MessageToDict(
|
|
155
|
+
cost,
|
|
156
|
+
preserving_proto_field_name=True,
|
|
157
|
+
always_print_fields_with_no_presence=True,
|
|
158
|
+
)
|
|
159
|
+
for cost in response.costs
|
|
160
|
+
]
|
|
161
|
+
logger.debug("Filtered costs retrieved with cost_dict: %s", cost_data_list)
|
|
162
|
+
return [CostData.model_validate(cost_data) for cost_data in cost_data_list]
|
{digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/filesystem/default_filesystem.py
RENAMED
|
@@ -18,14 +18,15 @@ logger = logging.getLogger(__name__)
|
|
|
18
18
|
class DefaultFilesystem(FilesystemStrategy):
|
|
19
19
|
"""Default state filesystem strategy."""
|
|
20
20
|
|
|
21
|
-
def __init__(self, mission_id: str, config: dict[str, str]) -> None:
|
|
21
|
+
def __init__(self, mission_id: str, setup_version_id: str, config: dict[str, str]) -> None:
|
|
22
22
|
"""Initialize the default filesystem strategy.
|
|
23
23
|
|
|
24
24
|
Args:
|
|
25
25
|
mission_id: The ID of the mission this strategy is associated with
|
|
26
|
+
setup_version_id: The ID of the setup version this strategy is associated with
|
|
26
27
|
config: A dictionary mapping names to Pydantic model classes
|
|
27
28
|
"""
|
|
28
|
-
super().__init__(mission_id, config)
|
|
29
|
+
super().__init__(mission_id, setup_version_id, config)
|
|
29
30
|
self.temp_root: str = self.config.get("temp_root", "") or tempfile.gettempdir()
|
|
30
31
|
os.makedirs(self.temp_root, exist_ok=True)
|
|
31
32
|
self.db: dict[str, FilesystemData] = {}
|
{digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/filesystem/filesystem_strategy.py
RENAMED
|
@@ -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
|
{digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/filesystem/grpc_filesystem.py
RENAMED
|
@@ -24,26 +24,6 @@ logger = logging.getLogger(__name__)
|
|
|
24
24
|
class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper):
|
|
25
25
|
"""Default state filesystem strategy."""
|
|
26
26
|
|
|
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
27
|
@staticmethod
|
|
48
28
|
@contextmanager
|
|
49
29
|
def _handle_grpc_errors(operation: str) -> Generator[Any, Any, Any]:
|
|
@@ -56,7 +36,7 @@ class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper):
|
|
|
56
36
|
operation: Description of the operation being performed.
|
|
57
37
|
|
|
58
38
|
Raises:
|
|
59
|
-
ValueError: Error
|
|
39
|
+
ValueError: Error with the model validation.
|
|
60
40
|
ServerError: from gRPC Client.
|
|
61
41
|
FilesystemServiceError: Filesystem service internal.
|
|
62
42
|
"""
|
|
@@ -71,6 +51,28 @@ class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper):
|
|
|
71
51
|
logger.exception(msg)
|
|
72
52
|
raise FilesystemServiceError(msg) from e
|
|
73
53
|
|
|
54
|
+
def __init__(
|
|
55
|
+
self,
|
|
56
|
+
mission_id: str,
|
|
57
|
+
setup_version_id: str,
|
|
58
|
+
config: dict[str, str],
|
|
59
|
+
client_config: ClientConfig,
|
|
60
|
+
**kwargs, # noqa: ANN003, ARG002
|
|
61
|
+
) -> None:
|
|
62
|
+
"""Initialize the default filesystem strategy.
|
|
63
|
+
|
|
64
|
+
Args:
|
|
65
|
+
mission_id: The ID of the mission this strategy is associated with
|
|
66
|
+
setup_version_id: The ID of the setup version this strategy is associated with
|
|
67
|
+
config: A dictionary mapping names to Pydantic model classes
|
|
68
|
+
client_config: The server configuration object
|
|
69
|
+
kwargs: other optional arguments to pass to the parent class constructor
|
|
70
|
+
"""
|
|
71
|
+
super().__init__(mission_id, setup_version_id, config)
|
|
72
|
+
channel = self._init_channel(client_config)
|
|
73
|
+
self.stub = filesystem_service_pb2_grpc.FilesystemServiceStub(channel)
|
|
74
|
+
logger.info("Channel client 'Filesystem' initialized succesfully")
|
|
75
|
+
|
|
74
76
|
def upload(self, content: bytes, name: str, file_type: FileType) -> FilesystemData:
|
|
75
77
|
"""Create a new file in the file system.
|
|
76
78
|
|
|
@@ -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=
|
|
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]:
|
|
@@ -209,12 +209,13 @@ class DefaultStorage(StorageStrategy):
|
|
|
209
209
|
def __init__(
|
|
210
210
|
self,
|
|
211
211
|
mission_id: str,
|
|
212
|
+
setup_version_id: str,
|
|
212
213
|
config: dict[str, type[BaseModel]],
|
|
213
214
|
storage_file_path: str = "local_storage",
|
|
214
215
|
**kwargs, # noqa: ANN003, ARG002
|
|
215
216
|
) -> None:
|
|
216
217
|
"""Initialize the storage."""
|
|
217
|
-
super().__init__(mission_id=mission_id, config=config)
|
|
218
|
+
super().__init__(mission_id=mission_id, setup_version_id=setup_version_id, config=config)
|
|
218
219
|
self.storage_file_path = f"{self.mission_id}_{storage_file_path}.json"
|
|
219
220
|
self.storage_file = Path(self.storage_file_path)
|
|
220
221
|
self.storage = self._load_from_file()
|
|
@@ -37,7 +37,7 @@ class GrpcStorage(StorageStrategy, GrpcClientWrapper):
|
|
|
37
37
|
dtype = DataType[raw["data_type"]]
|
|
38
38
|
payload = raw.get("data", {})
|
|
39
39
|
|
|
40
|
-
validated = self._validate_data(
|
|
40
|
+
validated = self._validate_data(coll, payload)
|
|
41
41
|
return StorageRecord(
|
|
42
42
|
mission_id=mission,
|
|
43
43
|
collection=coll,
|
|
@@ -185,12 +185,13 @@ class GrpcStorage(StorageStrategy, GrpcClientWrapper):
|
|
|
185
185
|
def __init__(
|
|
186
186
|
self,
|
|
187
187
|
mission_id: str,
|
|
188
|
+
setup_version_id: str,
|
|
188
189
|
config: dict[str, type[BaseModel]],
|
|
189
190
|
client_config: ClientConfig,
|
|
190
191
|
**kwargs, # noqa: ANN003, ARG002
|
|
191
192
|
) -> None:
|
|
192
193
|
"""Initialize the storage."""
|
|
193
|
-
super().__init__(mission_id=mission_id, config=config)
|
|
194
|
+
super().__init__(mission_id=mission_id, setup_version_id=setup_version_id, config=config)
|
|
194
195
|
|
|
195
196
|
channel = self._init_channel(client_config)
|
|
196
197
|
self.stub = storage_service_pb2_grpc.StorageServiceStub(channel)
|
|
@@ -39,11 +39,11 @@ class StorageRecord(BaseModel):
|
|
|
39
39
|
class StorageStrategy(BaseStrategy, ABC):
|
|
40
40
|
"""Define CRUD + list/remove-collection against a collection/record store."""
|
|
41
41
|
|
|
42
|
-
def _validate_data(self,
|
|
42
|
+
def _validate_data(self, collection: str, data: dict[str, Any]) -> BaseModel:
|
|
43
43
|
"""Validate data against the model schema for the given key.
|
|
44
44
|
|
|
45
45
|
Args:
|
|
46
|
-
|
|
46
|
+
collection: The unique name for the record type
|
|
47
47
|
data: The data to validate
|
|
48
48
|
|
|
49
49
|
Returns:
|
|
@@ -52,15 +52,15 @@ class StorageStrategy(BaseStrategy, ABC):
|
|
|
52
52
|
Raises:
|
|
53
53
|
ValueError: If the key has no associated model or validation fails
|
|
54
54
|
"""
|
|
55
|
-
model_cls = self.config.get(
|
|
55
|
+
model_cls = self.config.get(collection)
|
|
56
56
|
if not model_cls:
|
|
57
|
-
msg = f"No schema registered for collection '{
|
|
57
|
+
msg = f"No schema registered for collection '{collection}'"
|
|
58
58
|
raise ValueError(msg)
|
|
59
59
|
|
|
60
60
|
try:
|
|
61
61
|
return model_cls.model_validate(data)
|
|
62
62
|
except Exception as e:
|
|
63
|
-
msg = f"Validation failed for '{
|
|
63
|
+
msg = f"Validation failed for '{collection}': {e!s}"
|
|
64
64
|
raise ValueError(msg) from e
|
|
65
65
|
|
|
66
66
|
def _create_storage_record(
|
|
@@ -163,14 +163,15 @@ class StorageStrategy(BaseStrategy, ABC):
|
|
|
163
163
|
True if the deletion was successful, False otherwise
|
|
164
164
|
"""
|
|
165
165
|
|
|
166
|
-
def __init__(self, mission_id: str, config: dict[str, type[BaseModel]]) -> None:
|
|
166
|
+
def __init__(self, mission_id: str, setup_version_id: str, config: dict[str, type[BaseModel]]) -> None:
|
|
167
167
|
"""Initialize the storage strategy.
|
|
168
168
|
|
|
169
169
|
Args:
|
|
170
170
|
mission_id: The ID of the mission this strategy is associated with
|
|
171
|
+
setup_version_id: The ID of the setup version
|
|
171
172
|
config: A dictionary mapping names to Pydantic model classes
|
|
172
173
|
"""
|
|
173
|
-
super().__init__(mission_id)
|
|
174
|
+
super().__init__(mission_id, setup_version_id)
|
|
174
175
|
# Schema configuration mapping keys to model classes
|
|
175
176
|
self.config: dict[str, type[BaseModel]] = config
|
|
176
177
|
|
|
@@ -200,7 +201,7 @@ class StorageStrategy(BaseStrategy, ABC):
|
|
|
200
201
|
raise ValueError(msg)
|
|
201
202
|
record_id = record_id or uuid4().hex
|
|
202
203
|
data_type_enum = DataType[data_type]
|
|
203
|
-
validated_data = self._validate_data(
|
|
204
|
+
validated_data = self._validate_data(collection, {**data, "mission_id": self.mission_id})
|
|
204
205
|
record = self._create_storage_record(collection, record_id, validated_data, data_type_enum)
|
|
205
206
|
return self._store(record)
|
|
206
207
|
|
|
@@ -227,7 +228,7 @@ class StorageStrategy(BaseStrategy, ABC):
|
|
|
227
228
|
Returns:
|
|
228
229
|
StorageRecord: The modified record
|
|
229
230
|
"""
|
|
230
|
-
validated_data = self._validate_data(
|
|
231
|
+
validated_data = self._validate_data(collection, data)
|
|
231
232
|
return self._update(collection, record_id, validated_data)
|
|
232
233
|
|
|
233
234
|
def remove(self, collection: str, record_id: str) -> bool:
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: digitalkin
|
|
3
|
-
Version: 0.2.
|
|
3
|
+
Version: 0.2.12
|
|
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,7 +452,7 @@ Classifier: License :: Other/Proprietary License
|
|
|
452
452
|
Requires-Python: >=3.10
|
|
453
453
|
Description-Content-Type: text/markdown
|
|
454
454
|
License-File: LICENSE
|
|
455
|
-
Requires-Dist: digitalkin-proto>=0.1.
|
|
455
|
+
Requires-Dist: digitalkin-proto>=0.1.10
|
|
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
|
|
@@ -1,45 +0,0 @@
|
|
|
1
|
-
"""This module contains the abstract base class for cost strategies."""
|
|
2
|
-
|
|
3
|
-
from abc import ABC, abstractmethod
|
|
4
|
-
from enum import Enum, auto
|
|
5
|
-
from typing import Any
|
|
6
|
-
|
|
7
|
-
from pydantic import BaseModel
|
|
8
|
-
|
|
9
|
-
from digitalkin.services.base_strategy import BaseStrategy
|
|
10
|
-
|
|
11
|
-
|
|
12
|
-
class CostType(Enum):
|
|
13
|
-
"""."""
|
|
14
|
-
|
|
15
|
-
OTHER = auto()
|
|
16
|
-
TOKEN_INPUT = auto()
|
|
17
|
-
TOKEN_OUTPUT = auto()
|
|
18
|
-
API_CALL = auto()
|
|
19
|
-
STORAGE = auto()
|
|
20
|
-
TIME = auto()
|
|
21
|
-
|
|
22
|
-
|
|
23
|
-
class CostData(BaseModel):
|
|
24
|
-
"""."""
|
|
25
|
-
|
|
26
|
-
cost: float
|
|
27
|
-
mission_id: str
|
|
28
|
-
name: str
|
|
29
|
-
type: CostType
|
|
30
|
-
unit: str
|
|
31
|
-
|
|
32
|
-
|
|
33
|
-
class CostStrategy(BaseStrategy, ABC):
|
|
34
|
-
"""Abstract base class for cost strategies."""
|
|
35
|
-
|
|
36
|
-
@abstractmethod
|
|
37
|
-
def add_cost(self, cost_dict: dict[str, Any]) -> str:
|
|
38
|
-
"""Register a new cost."""
|
|
39
|
-
|
|
40
|
-
def __post_init__(self, *args, **kwargs) -> None: # noqa: ANN002, ANN003
|
|
41
|
-
"""Allow post init configuration."""
|
|
42
|
-
|
|
43
|
-
@abstractmethod
|
|
44
|
-
def get(self, cost_dict: dict[str, Any]) -> list[CostData]:
|
|
45
|
-
"""Get a cost."""
|
|
@@ -1,30 +0,0 @@
|
|
|
1
|
-
"""Default cost."""
|
|
2
|
-
|
|
3
|
-
import logging
|
|
4
|
-
from typing import Any
|
|
5
|
-
|
|
6
|
-
from digitalkin.services.cost.cost_strategy import CostData, CostStrategy
|
|
7
|
-
|
|
8
|
-
logger = logging.getLogger(__name__)
|
|
9
|
-
|
|
10
|
-
|
|
11
|
-
class DefaultCost(CostStrategy):
|
|
12
|
-
"""Default cost strategy."""
|
|
13
|
-
|
|
14
|
-
def add_cost(self, cost_dict: dict[str, Any]) -> str: # noqa: PLR6301
|
|
15
|
-
"""Create a new record in the cost database.
|
|
16
|
-
|
|
17
|
-
Returns:
|
|
18
|
-
str: The ID of the new record
|
|
19
|
-
"""
|
|
20
|
-
logger.info("Cost added with cost_dict: %s", cost_dict)
|
|
21
|
-
return ""
|
|
22
|
-
|
|
23
|
-
def get(self, cost_dict: dict[str, Any]) -> list[CostData]: # noqa: PLR6301
|
|
24
|
-
"""Get records from the database.
|
|
25
|
-
|
|
26
|
-
Returns:
|
|
27
|
-
list[CostData]: The list of records
|
|
28
|
-
"""
|
|
29
|
-
logger.info("Costs querried with cost_dict: %s", cost_dict)
|
|
30
|
-
return []
|
|
@@ -1,81 +0,0 @@
|
|
|
1
|
-
"""This module implements the default Cost strategy."""
|
|
2
|
-
|
|
3
|
-
import logging
|
|
4
|
-
from typing import Any
|
|
5
|
-
|
|
6
|
-
from digitalkin_proto.digitalkin.cost.v1 import cost_pb2, cost_service_pb2_grpc
|
|
7
|
-
from google.protobuf import json_format
|
|
8
|
-
from pydantic import ValidationError
|
|
9
|
-
|
|
10
|
-
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__)
|
|
15
|
-
|
|
16
|
-
|
|
17
|
-
class GrpcCost(CostStrategy, GrpcClientWrapper):
|
|
18
|
-
"""This class implements the default Cost strategy."""
|
|
19
|
-
|
|
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")
|
|
47
|
-
|
|
48
|
-
def add_cost(self, cost_dict: dict[str, Any]) -> str:
|
|
49
|
-
"""Create a new record in the cost database.
|
|
50
|
-
|
|
51
|
-
Required arguments:
|
|
52
|
-
data: Object representation of CostData
|
|
53
|
-
|
|
54
|
-
Returns:
|
|
55
|
-
str: The ID of the new record
|
|
56
|
-
"""
|
|
57
|
-
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 ""
|
|
65
|
-
|
|
66
|
-
request = cost_pb2.AddCostRequest(**valid_data.model_dump())
|
|
67
|
-
return self.exec_grpc_query("AddCost", request)
|
|
68
|
-
|
|
69
|
-
def get(self, cost_dict: dict[str, Any]) -> list[CostData]:
|
|
70
|
-
"""Get records from the database.
|
|
71
|
-
|
|
72
|
-
Returns:
|
|
73
|
-
list[CostData]: The list of records
|
|
74
|
-
"""
|
|
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)
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/grpc_servers/utils/grpc_client_wrapper.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/identity/default_identity.py
RENAMED
|
File without changes
|
{digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/identity/identity_strategy.py
RENAMED
|
File without changes
|
|
File without changes
|
{digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/registry/default_registry.py
RENAMED
|
File without changes
|
{digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/registry/registry_strategy.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/snapshot/default_snapshot.py
RENAMED
|
File without changes
|
{digitalkin-0.2.10 → digitalkin-0.2.12}/src/digitalkin/services/snapshot/snapshot_strategy.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|