digitalkin 0.3.2.dev3__py3-none-any.whl → 0.3.2.dev4__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.
digitalkin/__version__.py CHANGED
@@ -5,4 +5,4 @@ from importlib.metadata import PackageNotFoundError, version
5
5
  try:
6
6
  __version__ = version("digitalkin")
7
7
  except PackageNotFoundError:
8
- __version__ = "0.3.2.dev3"
8
+ __version__ = "0.3.2.dev4"
@@ -143,7 +143,8 @@ class SingleJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
143
143
  job_id: The unique identifier of the job.
144
144
  output_data: The output data produced by the job.
145
145
  """
146
- await self.tasks_sessions[job_id].queue.put(output_data.model_dump())
146
+ session = self.tasks_sessions[job_id]
147
+ await session.queue.put(output_data.model_dump())
147
148
 
148
149
  @asynccontextmanager # type: ignore
149
150
  async def generate_stream_consumer(self, job_id: str) -> AsyncIterator[AsyncGenerator[dict[str, Any], None]]: # type: ignore
@@ -262,6 +263,18 @@ class SingleJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
262
263
  logger.info("Managed task started: '%s'", job_id, extra={"task_id": job_id})
263
264
  return job_id
264
265
 
266
+ async def clean_session(self, task_id: str, mission_id: str) -> bool:
267
+ """Clean a task's session.
268
+
269
+ Args:
270
+ task_id: Unique identifier for the task.
271
+ mission_id: Mission identifier.
272
+
273
+ Returns:
274
+ bool: True if the task was successfully cleaned, False otherwise.
275
+ """
276
+ return await self._task_manager.clean_session(task_id, mission_id)
277
+
265
278
  async def stop_module(self, job_id: str) -> bool:
266
279
  """Stop a running module job.
267
280
 
@@ -287,7 +300,7 @@ class SingleJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
287
300
  await self.cancel_task(job_id, session.mission_id)
288
301
  logger.debug(
289
302
  "Module stopped successfully",
290
- extra={"job_id": job_id, "mission_id": session.mission_id, "module_name": session.module.name},
303
+ extra={"job_id": job_id, "mission_id": session.mission_id},
291
304
  )
292
305
  except Exception:
293
306
  logger.exception("Error stopping module", extra={"job_id": job_id})
@@ -41,6 +41,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
41
41
  args: Namespace
42
42
  setup: SetupStrategy
43
43
  job_manager: BaseJobManager
44
+ _registry_cache: RegistryStrategy | None = None
44
45
 
45
46
  def _add_parser_args(self, parser: ArgumentParser) -> None:
46
47
  super()._add_parser_args(parser)
@@ -84,11 +85,14 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
84
85
  self.setup = GrpcSetup() if self.args.services_mode == ServicesMode.REMOTE else DefaultSetup()
85
86
 
86
87
  def _get_registry(self) -> RegistryStrategy | None:
87
- """Get a registry instance if configured.
88
+ """Get a cached registry instance if configured.
88
89
 
89
90
  Returns:
90
- GrpcRegistry instance if registry config exists, None otherwise.
91
+ Cached GrpcRegistry instance if registry config exists, None otherwise.
91
92
  """
93
+ if self._registry_cache is not None:
94
+ return self._registry_cache
95
+
92
96
  registry_config = self.module_class.services_config_params.get("registry")
93
97
  if not registry_config:
94
98
  return None
@@ -97,7 +101,8 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
97
101
  if not client_config:
98
102
  return None
99
103
 
100
- return GrpcRegistry("", "", "", client_config)
104
+ self._registry_cache = GrpcRegistry("", "", "", client_config)
105
+ return self._registry_cache
101
106
 
102
107
  async def ConfigSetupModule( # noqa: N802
103
108
  self,
@@ -143,9 +148,13 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
143
148
  raise ServicerError(msg)
144
149
 
145
150
  # Resolve tool references in config_setup_data if registry is configured
151
+ # This also builds the tool_cache for LLM access during execution
146
152
  registry = self._get_registry()
147
153
  if registry:
148
- config_setup_data.resolve_tool_references(registry)
154
+ if hasattr(config_setup_data, "resolve_tool_references"):
155
+ config_setup_data.resolve_tool_references(registry)
156
+ if hasattr(config_setup_data, "build_tool_cache"):
157
+ config_setup_data.build_tool_cache()
149
158
 
150
159
  # create a task to run the module in background
151
160
  job_id = await self.job_manager.create_config_setup_instance_job(
@@ -6,6 +6,7 @@ from types import SimpleNamespace
6
6
  from typing import TYPE_CHECKING, Any
7
7
  from zoneinfo import ZoneInfo
8
8
 
9
+ from digitalkin.models.module.tool_cache import ToolCache
9
10
  from digitalkin.services.agent.agent_strategy import AgentStrategy
10
11
  from digitalkin.services.communication.communication_strategy import CommunicationStrategy
11
12
  from digitalkin.services.cost.cost_strategy import CostStrategy
@@ -101,7 +102,7 @@ class ModuleContext:
101
102
  metadata: SimpleNamespace
102
103
  helpers: SimpleNamespace
103
104
  state: SimpleNamespace = SimpleNamespace()
104
- tool_cache: dict[str, "ModuleInfo"]
105
+ tool_cache: ToolCache
105
106
 
106
107
  def __init__( # noqa: PLR0913, PLR0917
107
108
  self,
@@ -118,7 +119,7 @@ class ModuleContext:
118
119
  metadata: dict[str, Any] = {},
119
120
  helpers: dict[str, Any] = {},
120
121
  callbacks: dict[str, Any] = {},
121
- tool_cache: dict[str, "ModuleInfo"] | None = None,
122
+ tool_cache: ToolCache | None = None,
122
123
  ) -> None:
123
124
  """Register mandatory services, session, metadata and callbacks.
124
125
 
@@ -136,7 +137,7 @@ class ModuleContext:
136
137
  helpers: dict different user defined helpers.
137
138
  session: dict referring the session IDs or informations.
138
139
  callbacks: Functions allowing user to agent interaction.
139
- tool_cache: Pre-resolved tool references from setup.
140
+ tool_cache: ToolCache with pre-resolved tool references from setup.
140
141
  """
141
142
  # Core services
142
143
  self.agent = agent
@@ -153,15 +154,33 @@ class ModuleContext:
153
154
  self.session = Session(**session)
154
155
  self.helpers = SimpleNamespace(**helpers)
155
156
  self.callbacks = SimpleNamespace(**callbacks)
156
- self.tool_cache = tool_cache or {}
157
+ self.tool_cache = tool_cache or ToolCache()
157
158
 
158
- def get_tool(self, field_name: str) -> "ModuleInfo | None":
159
- """Get resolved tool info by setup field name.
159
+ def get_tool(self, slug: str) -> "ModuleInfo | None":
160
+ """Get resolved tool info by slug.
161
+
162
+ Fast lookup from the pre-populated tool cache.
163
+
164
+ Args:
165
+ slug: The tool slug to look up.
166
+
167
+ Returns:
168
+ ModuleInfo if found and valid, None otherwise.
169
+ """
170
+ return self.tool_cache.get(slug)
171
+
172
+ def check_and_get_tool(self, slug: str) -> "ModuleInfo | None":
173
+ """Check cache first, then query registry if not found.
174
+
175
+ This is the primary method for LLMs to discover tools. It:
176
+ 1. Checks the pre-populated cache (fast path)
177
+ 2. If not in cache, queries the registry
178
+ 3. If found via registry, caches the result
160
179
 
161
180
  Args:
162
- field_name: The name of the ToolReference field in the setup model.
181
+ slug: The tool slug to look up.
163
182
 
164
183
  Returns:
165
184
  ModuleInfo if found, None otherwise.
166
185
  """
167
- return self.tool_cache.get(field_name)
186
+ return self.tool_cache.check_and_get(slug, self.registry)
@@ -5,11 +5,12 @@ from __future__ import annotations
5
5
  import copy
6
6
  import types
7
7
  import typing
8
- from typing import TYPE_CHECKING, Any, Generic, TypeVar, cast, get_args, get_origin
8
+ from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar, cast, get_args, get_origin
9
9
 
10
- from pydantic import BaseModel, ConfigDict, create_model
10
+ from pydantic import BaseModel, ConfigDict, PrivateAttr, create_model
11
11
 
12
12
  from digitalkin.logger import logger
13
+ from digitalkin.models.module.tool_cache import ToolCache
13
14
  from digitalkin.models.module.tool_reference import ToolReference
14
15
  from digitalkin.utils.dynamic_schema import (
15
16
  DynamicField,
@@ -33,6 +34,9 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
33
34
  to be used to initialize the Kin. Supports dynamic schema providers for
34
35
  runtime value generation.
35
36
 
37
+ The tool_cache is populated during run_config_setup and contains resolved
38
+ ModuleInfo indexed by slug. It is validated during initialize.
39
+
36
40
  Attributes:
37
41
  model_fields: Inherited from Pydantic BaseModel, contains field definitions.
38
42
 
@@ -41,6 +45,9 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
41
45
  - Tests: tests/modules/test_setup_model.py
42
46
  """
43
47
 
48
+ _clean_model_cache: ClassVar[dict[tuple[type, bool, bool], type]] = {}
49
+ _tool_cache: ToolCache = PrivateAttr(default_factory=ToolCache)
50
+
44
51
  @classmethod
45
52
  async def get_clean_model(
46
53
  cls,
@@ -71,6 +78,11 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
71
78
  Returns:
72
79
  A new BaseModel subclass with filtered fields.
73
80
  """
81
+ # Check cache for non-forced requests
82
+ cache_key = (cls, config_fields, hidden_fields)
83
+ if not force and cache_key in cls._clean_model_cache:
84
+ return cast("type[SetupModelT]", cls._clean_model_cache[cache_key])
85
+
74
86
  clean_fields: dict[str, Any] = {}
75
87
 
76
88
  for name, field_info in cls.model_fields.items():
@@ -117,6 +129,11 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
117
129
  __config__=ConfigDict(arbitrary_types_allowed=True),
118
130
  **clean_fields,
119
131
  )
132
+
133
+ # Cache for non-forced requests
134
+ if not force:
135
+ cls._clean_model_cache[cache_key] = m
136
+
120
137
  return cast("type[SetupModelT]", m)
121
138
 
122
139
  @classmethod
@@ -461,3 +478,96 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
461
478
  cls._resolve_single_tool_reference("dict_value", item, registry)
462
479
  elif isinstance(item, BaseModel):
463
480
  cls._resolve_tool_references_recursive(item, registry)
481
+
482
+ @property
483
+ def tool_cache(self) -> ToolCache:
484
+ """Get the tool cache for this setup instance.
485
+
486
+ Returns:
487
+ The ToolCache containing resolved tools.
488
+ """
489
+ return self._tool_cache
490
+
491
+ def build_tool_cache(self) -> ToolCache:
492
+ """Build the tool cache from resolved ToolReferences.
493
+
494
+ This should be called during run_config_setup after resolve_tool_references.
495
+ It walks all ToolReference fields and adds resolved ones to the cache.
496
+
497
+ Returns:
498
+ The populated ToolCache.
499
+ """
500
+ self._build_tool_cache_recursive(self)
501
+ logger.debug(
502
+ "Tool cache built",
503
+ extra={"slugs": self._tool_cache.list_slugs()},
504
+ )
505
+ return self._tool_cache
506
+
507
+ def _build_tool_cache_recursive(self, model_instance: BaseModel) -> None:
508
+ """Recursively build tool cache from model fields.
509
+
510
+ Args:
511
+ model_instance: The model instance to process.
512
+ """
513
+ for field_name, field_value in model_instance.__dict__.items():
514
+ if field_value is None:
515
+ continue
516
+
517
+ if isinstance(field_value, ToolReference):
518
+ self._add_tool_reference_to_cache(field_name, field_value)
519
+ elif isinstance(field_value, BaseModel):
520
+ self._build_tool_cache_recursive(field_value)
521
+ elif isinstance(field_value, list):
522
+ self._build_tool_cache_from_list(field_value)
523
+ elif isinstance(field_value, dict):
524
+ self._build_tool_cache_from_dict(field_value)
525
+
526
+ def _add_tool_reference_to_cache(self, field_name: str, tool_ref: ToolReference) -> None:
527
+ """Add a resolved ToolReference to the cache.
528
+
529
+ Args:
530
+ field_name: Name of the field (used as fallback slug).
531
+ tool_ref: The ToolReference instance.
532
+ """
533
+ if tool_ref.module_info:
534
+ # Use slug from config, or field name as fallback
535
+ slug = tool_ref.slug or field_name
536
+ self._tool_cache.add(slug, tool_ref.module_info)
537
+
538
+ def _build_tool_cache_from_list(self, items: list) -> None:
539
+ """Build tool cache from list items.
540
+
541
+ Args:
542
+ items: List of items to process.
543
+ """
544
+ for idx, item in enumerate(items):
545
+ if isinstance(item, ToolReference):
546
+ self._add_tool_reference_to_cache(f"list_{idx}", item)
547
+ elif isinstance(item, BaseModel):
548
+ self._build_tool_cache_recursive(item)
549
+
550
+ def _build_tool_cache_from_dict(self, mapping: dict) -> None:
551
+ """Build tool cache from dict values.
552
+
553
+ Args:
554
+ mapping: Dict to process.
555
+ """
556
+ for key, item in mapping.items():
557
+ if isinstance(item, ToolReference):
558
+ self._add_tool_reference_to_cache(str(key), item)
559
+ elif isinstance(item, BaseModel):
560
+ self._build_tool_cache_recursive(item)
561
+
562
+ def validate_tool_cache(self, registry: RegistryStrategy) -> list[str]:
563
+ """Validate all cached tools are still available.
564
+
565
+ Should be called during initialize to ensure tools are still valid.
566
+
567
+ Args:
568
+ registry: Registry to validate against.
569
+
570
+ Returns:
571
+ List of slugs that are no longer valid.
572
+ """
573
+ return self._tool_cache.validate(registry)
@@ -0,0 +1,230 @@
1
+ """Tool cache for managing resolved tool references.
2
+
3
+ The ToolCache is a registry that stores resolved ModuleInfo by slug.
4
+ It is populated during run_config_setup and validated during initialize.
5
+ LLMs check the cache before calling the registry for tool discovery.
6
+ """
7
+
8
+ from __future__ import annotations
9
+
10
+ from typing import TYPE_CHECKING
11
+
12
+ from pydantic import BaseModel, Field
13
+
14
+ from digitalkin.logger import logger
15
+
16
+ if TYPE_CHECKING:
17
+ from digitalkin.models.services.registry import ModuleInfo
18
+ from digitalkin.services.registry import RegistryStrategy
19
+
20
+
21
+ class ToolCacheEntry(BaseModel):
22
+ """Single entry in the tool cache."""
23
+
24
+ slug: str = Field(description="Unique identifier/slug for this tool")
25
+ module_id: str = Field(description="Resolved module ID")
26
+ module_info: ModuleInfo = Field(description="Full module information")
27
+ is_valid: bool = Field(default=True, description="Whether this entry is still valid")
28
+
29
+ model_config = {"arbitrary_types_allowed": True}
30
+
31
+
32
+ class ToolCache(BaseModel):
33
+ """Registry cache for resolved tools.
34
+
35
+ Stores tool references by slug during run_config_setup and provides
36
+ lookup methods for LLMs during execution. Tools must be checked
37
+ against the cache before calling the registry.
38
+
39
+ Flow:
40
+ 1. run_config_setup: Registry called, results cached by slug
41
+ 2. initialize: Cache validated (tools still available)
42
+ 3. Runtime: LLM checks cache first, then registry if not found
43
+ """
44
+
45
+ entries: dict[str, ToolCacheEntry] = Field(default_factory=dict)
46
+
47
+ def add(self, slug: str, module_info: ModuleInfo) -> None:
48
+ """Add a tool to the cache.
49
+
50
+ Args:
51
+ slug: Unique identifier for this tool.
52
+ module_info: Resolved module information.
53
+ """
54
+ self.entries[slug] = ToolCacheEntry(
55
+ slug=slug,
56
+ module_id=module_info.module_id,
57
+ module_info=module_info,
58
+ is_valid=True,
59
+ )
60
+ logger.debug("Tool cached", extra={"slug": slug, "module_id": module_info.module_id})
61
+
62
+ def get(self, slug: str) -> ModuleInfo | None:
63
+ """Get a tool from the cache by slug.
64
+
65
+ Args:
66
+ slug: The tool slug to look up.
67
+
68
+ Returns:
69
+ ModuleInfo if found and valid, None otherwise.
70
+ """
71
+ entry = self.entries.get(slug)
72
+ if entry and entry.is_valid:
73
+ return entry.module_info
74
+ return None
75
+
76
+ def contains(self, slug: str) -> bool:
77
+ """Check if a tool exists in the cache.
78
+
79
+ Args:
80
+ slug: The tool slug to check.
81
+
82
+ Returns:
83
+ True if tool exists and is valid.
84
+ """
85
+ entry = self.entries.get(slug)
86
+ return entry is not None and entry.is_valid
87
+
88
+ def invalidate(self, slug: str) -> None:
89
+ """Mark a tool as invalid.
90
+
91
+ Args:
92
+ slug: The tool slug to invalidate.
93
+ """
94
+ if slug in self.entries:
95
+ self.entries[slug].is_valid = False
96
+ logger.debug("Tool invalidated", extra={"slug": slug})
97
+
98
+ def remove(self, slug: str) -> None:
99
+ """Remove a tool from the cache.
100
+
101
+ Args:
102
+ slug: The tool slug to remove.
103
+ """
104
+ if slug in self.entries:
105
+ del self.entries[slug]
106
+ logger.debug("Tool removed from cache", extra={"slug": slug})
107
+
108
+ def clear(self) -> None:
109
+ """Clear all entries from the cache."""
110
+ self.entries.clear()
111
+
112
+ def check_and_get(
113
+ self,
114
+ slug: str,
115
+ registry: RegistryStrategy | None = None,
116
+ ) -> ModuleInfo | None:
117
+ """Check cache first, then optionally query registry.
118
+
119
+ This is the primary method for LLMs to discover tools. It:
120
+ 1. Checks if the tool is in cache (fast path)
121
+ 2. If not in cache and registry provided, queries registry
122
+ 3. If found via registry, caches the result
123
+
124
+ Args:
125
+ slug: The tool slug to look up.
126
+ registry: Optional registry to query if not in cache.
127
+
128
+ Returns:
129
+ ModuleInfo if found, None otherwise.
130
+ """
131
+ # Fast path: check cache
132
+ cached = self.get(slug)
133
+ if cached:
134
+ logger.debug("Tool cache hit", extra={"slug": slug})
135
+ return cached
136
+
137
+ # Not in cache - try registry if available
138
+ if registry:
139
+ logger.debug("Tool cache miss, querying registry", extra={"slug": slug})
140
+ try:
141
+ # Try by ID first (slug might be module_id)
142
+ info = registry.discover_by_id(slug)
143
+ if info:
144
+ self.add(slug, info)
145
+ return info
146
+
147
+ # Try by tag search
148
+ results = registry.search(name=slug, module_type="tool", organization_id=None)
149
+ if results:
150
+ info = results[0]
151
+ self.add(slug, info)
152
+ return info
153
+ except Exception:
154
+ logger.exception("Registry lookup failed", extra={"slug": slug})
155
+
156
+ return None
157
+
158
+ def validate(self, registry: RegistryStrategy) -> list[str]:
159
+ """Validate all cached tools are still available.
160
+
161
+ Checks each cached tool against the registry and marks
162
+ invalid entries. Returns list of invalid slugs.
163
+
164
+ Args:
165
+ registry: Registry to validate against.
166
+
167
+ Returns:
168
+ List of slugs that are no longer valid.
169
+ """
170
+ invalid: list[str] = []
171
+ for slug, entry in self.entries.items():
172
+ if not entry.is_valid:
173
+ continue
174
+ try:
175
+ info = registry.discover_by_id(entry.module_id)
176
+ if not info:
177
+ entry.is_valid = False
178
+ invalid.append(slug)
179
+ logger.warning("Tool no longer available", extra={"slug": slug, "module_id": entry.module_id})
180
+ except Exception:
181
+ entry.is_valid = False
182
+ invalid.append(slug)
183
+ logger.exception("Tool validation failed", extra={"slug": slug})
184
+ return invalid
185
+
186
+ def list_slugs(self) -> list[str]:
187
+ """Get list of all valid tool slugs.
188
+
189
+ Returns:
190
+ List of valid tool slugs.
191
+ """
192
+ return [slug for slug, entry in self.entries.items() if entry.is_valid]
193
+
194
+ def to_dict(self) -> dict[str, dict]:
195
+ """Serialize cache to dict for storage.
196
+
197
+ Returns:
198
+ Dict representation of cache entries.
199
+ """
200
+ return {
201
+ slug: {
202
+ "slug": entry.slug,
203
+ "module_id": entry.module_id,
204
+ "is_valid": entry.is_valid,
205
+ "module_info": entry.module_info.model_dump(),
206
+ }
207
+ for slug, entry in self.entries.items()
208
+ }
209
+
210
+ @classmethod
211
+ def from_dict(cls, data: dict[str, dict]) -> ToolCache:
212
+ """Deserialize cache from dict.
213
+
214
+ Args:
215
+ data: Dict representation of cache entries.
216
+
217
+ Returns:
218
+ ToolCache instance.
219
+ """
220
+ from digitalkin.models.services.registry import ModuleInfo # noqa: PLC0415
221
+
222
+ cache = cls()
223
+ for slug, entry_data in data.items():
224
+ cache.entries[slug] = ToolCacheEntry(
225
+ slug=entry_data["slug"],
226
+ module_id=entry_data["module_id"],
227
+ module_info=ModuleInfo(**entry_data["module_info"]),
228
+ is_valid=entry_data.get("is_valid", True),
229
+ )
230
+ return cache
@@ -23,6 +23,7 @@ class ToolReferenceConfig(BaseModel):
23
23
  """Configuration for how a tool should be selected."""
24
24
 
25
25
  mode: ToolSelectionMode = Field(default=ToolSelectionMode.FIXED)
26
+ slug: str | None = Field(default=None, description="Unique slug for cache lookup")
26
27
  fixed_id: str | None = Field(default=None, description="Module ID for FIXED mode")
27
28
  tag: str | None = Field(default=None, description="Search tag for TAG mode")
28
29
  organization_id: str | None = Field(default=None, description="Filter by organization")
@@ -58,6 +59,20 @@ class ToolReference(BaseModel):
58
59
  selected_module_id: str | None = Field(default=None, description="Resolved module ID after resolution")
59
60
  _cached_info: ModuleInfo | None = PrivateAttr(default=None)
60
61
 
62
+ @property
63
+ def slug(self) -> str | None:
64
+ """Get the slug for cache lookup.
65
+
66
+ Returns config.slug if set, otherwise falls back to fixed_id or tag.
67
+ """
68
+ if self.config.slug:
69
+ return self.config.slug
70
+ if self.config.mode == ToolSelectionMode.FIXED and self.config.fixed_id:
71
+ return self.config.fixed_id
72
+ if self.config.mode == ToolSelectionMode.TAG and self.config.tag:
73
+ return self.config.tag
74
+ return None
75
+
61
76
  @property
62
77
  def module_info(self) -> ModuleInfo | None:
63
78
  """Get cached ModuleInfo if resolved."""
@@ -18,9 +18,7 @@ from digitalkin.models.module.module_types import (
18
18
  SecretModelT,
19
19
  SetupModelT,
20
20
  )
21
- from digitalkin.models.module.tool_reference import ToolReference
22
21
  from digitalkin.models.module.utility import EndOfStreamOutput, ModuleStartInfoOutput
23
- from digitalkin.models.services.registry import ModuleInfo
24
22
  from digitalkin.models.services.storage import BaseRole
25
23
  from digitalkin.modules.trigger_handler import TriggerHandler
26
24
  from digitalkin.services.services_config import ServicesConfig, ServicesStrategy
@@ -439,22 +437,6 @@ class BaseModule( # noqa: PLR0904
439
437
  else:
440
438
  self._status = ModuleStatus.STOPPING
441
439
 
442
- @staticmethod
443
- def _extract_tool_cache(setup_data: SetupModelT) -> dict[str, ModuleInfo]:
444
- """Extract resolved ToolReference info from setup data.
445
-
446
- Args:
447
- setup_data: The setup model containing potential ToolReference fields.
448
-
449
- Returns:
450
- Dict mapping field names to resolved ModuleInfo.
451
- """
452
- cache: dict[str, ModuleInfo] = {}
453
- for name, value in setup_data.__dict__.items():
454
- if isinstance(value, ToolReference) and value.module_info:
455
- cache[name] = value.module_info
456
- return cache
457
-
458
440
  async def start(
459
441
  self,
460
442
  input_data: InputModelT,
@@ -466,8 +448,19 @@ class BaseModule( # noqa: PLR0904
466
448
  try:
467
449
  self.context.callbacks.send_message = callback
468
450
 
469
- # Populate tool cache from resolved ToolReference fields
470
- self.context.tool_cache = self._extract_tool_cache(setup_data)
451
+ # Use tool cache from setup_data if available (populated during run_config_setup)
452
+ # Tool cache is optional - modules can work without it
453
+ if hasattr(setup_data, "tool_cache") and setup_data.tool_cache is not None:
454
+ self.context.tool_cache = setup_data.tool_cache
455
+
456
+ # Validate tool cache - ensure all cached tools are still available
457
+ if self.context.registry and hasattr(setup_data, "validate_tool_cache"):
458
+ invalid_tools = setup_data.validate_tool_cache(self.context.registry)
459
+ if invalid_tools:
460
+ logger.warning(
461
+ "Some cached tools are no longer available",
462
+ extra={"invalid_tools": invalid_tools},
463
+ )
471
464
 
472
465
  # Send module start info as first message
473
466
  await callback(
@@ -4,7 +4,6 @@ from typing import Any, ClassVar
4
4
 
5
5
  from pydantic import BaseModel, Field, PrivateAttr
6
6
 
7
- from digitalkin import logger
8
7
  from digitalkin.services.agent import AgentStrategy, DefaultAgent
9
8
  from digitalkin.services.communication import CommunicationStrategy, DefaultCommunication, GrpcCommunication
10
9
  from digitalkin.services.cost import CostStrategy, DefaultCost, GrpcCost
@@ -144,9 +143,6 @@ class ServicesConfig(BaseModel):
144
143
  msg = f"Strategy {name} not found in ServicesConfig."
145
144
  raise ValueError(msg)
146
145
 
147
- logger.logger.critical(
148
- f"{name=}\n{mission_id=}, {setup_id=}, {setup_version_id=}, {self.get_strategy_config(name)=}"
149
- )
150
146
  # Instantiate the strategy with the mission ID, setup version ID, and configuration
151
147
  return strategy_type(mission_id, setup_id, setup_version_id, **self.get_strategy_config(name) or {})
152
148
 
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: digitalkin
3
- Version: 0.3.2.dev3
3
+ Version: 0.3.2.dev4
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
@@ -7,7 +7,7 @@ base_server/mock/__init__.py,sha256=YZFT-F1l_TpvJYuIPX-7kTeE1CfOjhx9YmNRXVoi-jQ,
7
7
  base_server/mock/mock_pb2.py,sha256=sETakcS3PAAm4E-hTCV1jIVaQTPEAIoVVHupB8Z_k7Y,1843
8
8
  base_server/mock/mock_pb2_grpc.py,sha256=BbOT70H6q3laKgkHfOx1QdfmCS_HxCY4wCOX84YAdG4,3180
9
9
  digitalkin/__init__.py,sha256=7LLBAba0th-3SGqcpqFO-lopWdUkVLKzLZiMtB-mW3M,162
10
- digitalkin/__version__.py,sha256=eJ__RY2lPoQprF2CURDjtvLHzHcyg0TSiE4skNzcSoQ,195
10
+ digitalkin/__version__.py,sha256=k7ywgivs-7wcYmxpZLVkbYfyQZo7NuQeNm6D2K0VbkI,195
11
11
  digitalkin/logger.py,sha256=8ze_tjt2G6mDTuQcsf7-UTXWP3UHZ7LZVSs_iqF4rX4,4685
12
12
  digitalkin/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
13
13
  digitalkin/core/__init__.py,sha256=FJRcJ-B1Viyn-38L8XpOpZ8KOnf1I7PCDOAmKXLQhqc,71
@@ -15,7 +15,7 @@ digitalkin/core/common/__init__.py,sha256=Gh2eJAJRnrUE93jSEfG7r0nb01Xh1kSkNL6nEp
15
15
  digitalkin/core/common/factories.py,sha256=mV6SmXXrZxzIQ7DLdDtPdjapSDZt5Ua-nBIDsozs_Vk,5047
16
16
  digitalkin/core/job_manager/__init__.py,sha256=gGtgQpE6vbBHxAj1SYMbcpj45Q6x8IcsqnyQPfyZZ-8,25
17
17
  digitalkin/core/job_manager/base_job_manager.py,sha256=fRnoR74uduImGYOKj4kZ5bo0xv9hrXSCUuq0NHfDpPE,10089
18
- digitalkin/core/job_manager/single_job_manager.py,sha256=ivR9L6sUMdQ7i3bhIgPLESxgWWGl6IlaBz__4csIEro,14069
18
+ digitalkin/core/job_manager/single_job_manager.py,sha256=xOgl8HNUpfd_OPmoJpUt8hJznxDTeUzFpEFBRgEYzp4,14464
19
19
  digitalkin/core/job_manager/taskiq_broker.py,sha256=ORGg2QSLDaWiCqie8ZGkbcDq7LDuWWSx7hg-SoUv99E,11544
20
20
  digitalkin/core/job_manager/taskiq_job_manager.py,sha256=d23iylsc2RUvoE95VxZWa5Kb_yqptMOF63l7MDtuZtk,21264
21
21
  digitalkin/core/task_manager/__init__.py,sha256=k9i-qIoee_1yXogyQolaVFDUQBIZU3ENbYKtjrCNmTQ,31
@@ -28,7 +28,7 @@ digitalkin/core/task_manager/task_session.py,sha256=5jw21bT_SPXUzWE7tk6YG62EXqlR
28
28
  digitalkin/grpc_servers/__init__.py,sha256=ZIRMJ1Lcas8yQ106GCup6hn2UBOsx1sNk8ap0lpEDnY,72
29
29
  digitalkin/grpc_servers/_base_server.py,sha256=ZVeCDwI7w7fFbPTXPkeJb_SOuLfd2T7za3T4oCu2UWY,18680
30
30
  digitalkin/grpc_servers/module_server.py,sha256=7jO6Bz4pxno0AjRFercZMwgOFW7dutywynoCCht3dnQ,7858
31
- digitalkin/grpc_servers/module_servicer.py,sha256=DvoBEGzDgzWQq5PHt_ZqUgyl2NK7vPaSrOkfQPRITuY,20619
31
+ digitalkin/grpc_servers/module_servicer.py,sha256=a5qVOd0beAwXlFAch9VfLNR8PdaYj_fF2O8rXcwInSU,21087
32
32
  digitalkin/grpc_servers/utils/__init__.py,sha256=ZnAIb_F8z4NhtPypqkdmzgRSzolKnJTk3oZx5GfWH5Y,38
33
33
  digitalkin/grpc_servers/utils/exceptions.py,sha256=LtaDtlqXCeT6iqApogs4pbtezotOVeg4fhnFzGBvFsY,692
34
34
  digitalkin/grpc_servers/utils/grpc_client_wrapper.py,sha256=ElGvp6evY5q-EBmDVyQZaDJktfShtMsptMmq16jfgxA,3285
@@ -53,17 +53,18 @@ digitalkin/models/grpc_servers/types.py,sha256=rQ78s4nAet2jy-NIDj_PUWriT0kuGHr_w
53
53
  digitalkin/models/module/__init__.py,sha256=N55wan3rAUVPEGLIDjXoFM_-DYY_zxvbQOZHzNDfwoY,751
54
54
  digitalkin/models/module/base_types.py,sha256=oIylVNqo0idTFj4dRgCt7P19daNZ-AlvgCPpL9TJvto,1850
55
55
  digitalkin/models/module/module.py,sha256=k0W8vfJJFth8XdDzkHm32SyTuSf3h2qF0hSrxAfGF1s,956
56
- digitalkin/models/module/module_context.py,sha256=HDtKbQxwJU_RSE7xTA3kFkjXGQiDbG64q45fNudKUYI,5734
56
+ digitalkin/models/module/module_context.py,sha256=cnk5OpgevgqkmD_gU5KWSBzZ2vsNFUyN8rEzW0TNI-s,6354
57
57
  digitalkin/models/module/module_types.py,sha256=C9azCNBk76xMa-Mww8_6AiwQR8MLAsEyUOvBYxytovI,739
58
- digitalkin/models/module/setup_types.py,sha256=BzRLpYDAPbKLUOZ9C8ChTguuxzTveYYf4pC_0gda9a0,17215
59
- digitalkin/models/module/tool_reference.py,sha256=vg5isqXXouY4GFNS0mOodenojA_AWe1lISS6LABdEY8,3719
58
+ digitalkin/models/module/setup_types.py,sha256=e9CG8x75Mkh7dDIQ66UWU99VcWmyk-JAqHCpUtcgyns,21323
59
+ digitalkin/models/module/tool_cache.py,sha256=H8k2ue9m61oNHGMDgF9i3e4FxfFQoNW1iIZqoWyZNOc,7523
60
+ digitalkin/models/module/tool_reference.py,sha256=A4EewKMlfhdI2hW7F3kgPc6_UinUrN-hhOweM0KzSa0,4302
60
61
  digitalkin/models/module/utility.py,sha256=gnbYfWpXGbomUI0fWf7T-Qm_VvT-LXDv1OuA9zObwVg,5589
61
62
  digitalkin/models/services/__init__.py,sha256=jhfVw6egq0OcHmos_fypH9XFehbHTBw09wluVFVFEyw,226
62
63
  digitalkin/models/services/cost.py,sha256=9PXvd5RrIk9vCrRjcUGQ9ZyAokEbwLg4s0RfnE-aLP4,1616
63
64
  digitalkin/models/services/registry.py,sha256=0qMAIi52ovofZ4XTV7UkI4i7sR5C3_VFLp-kw9AjQ24,829
64
65
  digitalkin/models/services/storage.py,sha256=wp7F-AvTsU46ujGPcguqM5kUKRZx4399D4EGAAJt2zs,1143
65
66
  digitalkin/modules/__init__.py,sha256=vTQk8DWopxQSJ17BjE5dNhq247Rou55iQLJdBxoPUmo,296
66
- digitalkin/modules/_base_module.py,sha256=WLByYSOdXhwEzO-7S1ujIOzRa4TE1ckFQdjySPjaHeA,21321
67
+ digitalkin/modules/_base_module.py,sha256=0bbiPOcLSiKqjrVYaWdmc72jbpaPa8dOzozEcQO0zMg,21285
67
68
  digitalkin/modules/archetype_module.py,sha256=XC9tl1Yr6QlbPn_x0eov6UUZwQgwW--BYPPMYVJH_NU,505
68
69
  digitalkin/modules/tool_module.py,sha256=GBis7bKCkvWFCYLRvaS9oZVmLBBve1w8BhVnKOU2sCc,506
69
70
  digitalkin/modules/trigger_handler.py,sha256=qPNMi-8NHqscOxciHeaXtpwjXApT3YzjMF23zQAjaZY,1770
@@ -73,7 +74,7 @@ digitalkin/modules/triggers/healthcheck_services_trigger.py,sha256=TpPw5XTnw3Bt9
73
74
  digitalkin/modules/triggers/healthcheck_status_trigger.py,sha256=rozWQWvO7a2ZTg8BjFCyEWeAai5vdbi7BBwu0vR5jv8,1768
74
75
  digitalkin/services/__init__.py,sha256=DugdZZ3uKXFsA5C-DuP41ssGegU1FybI3HQcW_Se4bk,1101
75
76
  digitalkin/services/base_strategy.py,sha256=yA9KUJGRKuuaxA6l3GcMv8zKfWoIsW03UxJT80Yea2I,766
76
- digitalkin/services/services_config.py,sha256=PySj8qKJISFUMNmsSWM-1VmqQMCqyi1sIz5SX4GtZm8,8868
77
+ digitalkin/services/services_config.py,sha256=rz1cJdO_dfUN3kwOQ40XRSE4ZH3mgM43vXka94gjymY,8689
77
78
  digitalkin/services/services_models.py,sha256=Cb-ajJxHlEOQmoylMCfNdN08T6IOmMLXHrAZ6pxprrQ,1855
78
79
  digitalkin/services/agent/__init__.py,sha256=vJc8JN0pdtA8ecypLBeHrwAUIW6H2C8NyW-dk24rTpk,244
79
80
  digitalkin/services/agent/agent_strategy.py,sha256=42Q9RciHX6tg3CgDQkbrlIx4h_TX0WIuSpLmCjitVmA,492
@@ -120,7 +121,7 @@ digitalkin/utils/development_mode_action.py,sha256=2hznh0ajW_4ZTysfoc0Y49161f_PQ
120
121
  digitalkin/utils/dynamic_schema.py,sha256=5-B3dBGlCYYv6uRJkgudtc0ZpBOTYxl0yKedDGsteZQ,15184
121
122
  digitalkin/utils/llm_ready_schema.py,sha256=JjMug_lrQllqFoanaC091VgOqwAd-_YzcpqFlS7p778,2375
122
123
  digitalkin/utils/package_discover.py,sha256=sa6Zp5Kape1Zr4iYiNrnZxiHDnqM06ODk6yfWHom53w,13465
123
- digitalkin-0.3.2.dev3.dist-info/licenses/LICENSE,sha256=Ies4HFv2r2hzDRakJYxk3Y60uDFLiG-orIgeTpstnIo,20327
124
+ digitalkin-0.3.2.dev4.dist-info/licenses/LICENSE,sha256=Ies4HFv2r2hzDRakJYxk3Y60uDFLiG-orIgeTpstnIo,20327
124
125
  modules/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
125
126
  modules/cpu_intensive_module.py,sha256=GZlirQDZdYuXrI46sv1q4RNAHZjL4EptHVQTvgK9zz8,8363
126
127
  modules/dynamic_setup_module.py,sha256=tKvUWZdlYZkfAgKR0mLuFcLiFGKpVgpsz10LeJ6B2QI,11410
@@ -128,7 +129,7 @@ modules/minimal_llm_module.py,sha256=N9aIzZQI-miyH4AB4xTmGHpMvdSLnYyXNOD4Z3YFzis
128
129
  modules/text_transform_module.py,sha256=MfhI_Ki1U6qk379ne6oazNEu4PhO4R3cRezEcr0nGPw,7251
129
130
  services/filesystem_module.py,sha256=U4dgqtuDadaXz8PJ1d_uQ_1EPncBqudAQCLUICF9yL4,7421
130
131
  services/storage_module.py,sha256=Wz2MzLvqs2D_bnBBgtnujYcAKK2V2KFMk8K21RoepSE,6972
131
- digitalkin-0.3.2.dev3.dist-info/METADATA,sha256=-yHoriRYCcXVXhVrO_IOa3S3yRK0yCkZXhpfqP3S1KI,29724
132
- digitalkin-0.3.2.dev3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
133
- digitalkin-0.3.2.dev3.dist-info/top_level.txt,sha256=gcjqlyrZuLjIyxrOIavCQM_olpr6ND5kPKkZd2j0xGo,40
134
- digitalkin-0.3.2.dev3.dist-info/RECORD,,
132
+ digitalkin-0.3.2.dev4.dist-info/METADATA,sha256=9owo2Dtpnwbm4Rjkx9iXA_rKkzstywsbhe8hU9Mn8pY,29724
133
+ digitalkin-0.3.2.dev4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
134
+ digitalkin-0.3.2.dev4.dist-info/top_level.txt,sha256=gcjqlyrZuLjIyxrOIavCQM_olpr6ND5kPKkZd2j0xGo,40
135
+ digitalkin-0.3.2.dev4.dist-info/RECORD,,