digitalkin 1.0.0.dev10__tar.gz → 1.0.0.dev12__tar.gz

{digitalkin-1.0.0.dev10 → digitalkin-1.0.0.dev12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: digitalkin
-Version: 1.0.0.dev10
+Version: 1.0.0.dev12
 Summary: SDK to build kin used in DigitalKin
 Author-email: "DigitalKin.ai" <contact@digitalkin.ai>
 License: Attribution-NonCommercial-ShareAlike 4.0 International

{digitalkin-1.0.0.dev10 → digitalkin-1.0.0.dev12}/pyproject.toml

@@ -37,7 +37,7 @@
         "pydantic-settings>=2.14.1",
         "redis[hiredis]>=7.4.0,<8",
     ]
-    version = "1.0.0.dev10"
+    version = "1.0.0.dev12"
 
     [project.optional-dependencies]
         performance = [ "uvloop>=0.21" ]
@@ -274,6 +274,8 @@
 [tool.mypy]
     exclude                = [ "examples", "tests", "scripts" ]
     ignore_missing_imports = true
+    warn_unused_ignores    = true
+    warn_redundant_casts   = true
 
 [tool.pytest.ini_options]
 
@@ -294,6 +296,7 @@
         "contract: marks gRPC proto contract verification tests",
         "e2e: marks end-to-end tests requiring full Docker Compose stack",
         "edge_case: marks tests for boundary conditions and edge cases",
+        "flaky: marks known-flaky tests for the quarantine plugin (tests/fixtures/flakiness.py)",
         "grpc: marks tests for gRPC service functionality",
         "idempotency: marks tests for retry and duplicate request handling",
         "integration: marks tests that require external service connections (deselect with '-m \"not integration\"')",

{digitalkin-1.0.0.dev10 → digitalkin-1.0.0.dev12}/src/digitalkin/__version__.py

@@ -5,4 +5,4 @@ from importlib.metadata import PackageNotFoundError, version
 try:
     __version__ = version("digitalkin")
 except PackageNotFoundError:
-    __version__ = "1.0.0.dev10"
+    __version__ = "1.0.0.dev12"

{digitalkin-1.0.0.dev10 → digitalkin-1.0.0.dev12}/src/digitalkin/community/agno/__init__.py

@@ -4,28 +4,24 @@ Adapters, converters, and HITL helpers for building DigitalKin modules
 on top of the Agno agent framework. Exports:
 
 - :class:`AgnoStreamAdapter` — Agno streaming events → DigitalKin events.
-- :func:`agui_tool_to_external_function` / :func:`make_tools_factory` —
-  register AG-UI client-side (frontend) tools as Agno external Functions.
+- :class:`AguiTools` — register AG-UI client-side (frontend) tools as Agno
+  external Functions.
 - :class:`AgnoHitlRunner`, :class:`PausedRunStore`, :class:`PauseInfo`,
   :class:`PausedRunRecord`, :data:`HITL_STORAGE_CONFIG`,
-  :func:`emit_awaiting_tool_result` — human-in-the-loop (HITL) runner
-  that persists a paused Agno run via the module's
+  :class:`HitlEvents` — human-in-the-loop (HITL) runner that persists a
+  paused Agno run via the module's
   :class:`~digitalkin.services.storage.StorageStrategy` and resumes it
   when the front replies with a ``ToolMessage``.
 """
 
 from digitalkin.community.agno.agno_adapter import AgnoStreamAdapter
-from digitalkin.community.agno.agui_tools import (
-    agui_tool_to_external_function,
-    make_tools_factory,
-)
+from digitalkin.community.agno.agui_tools import AguiTools
 from digitalkin.community.agno.hitl import (
     HITL_STORAGE_CONFIG,
     AgnoHitlRunner,
+    HitlEvents,
     PausedRunRecord,
     PausedRunStore,
-    emit_awaiting_tool_result,
-    emit_messages_snapshot,
 )
 from digitalkin.community.agno.models import PauseInfo
 
@@ -33,11 +29,9 @@ __all__ = [
     "HITL_STORAGE_CONFIG",
     "AgnoHitlRunner",
     "AgnoStreamAdapter",
+    "AguiTools",
+    "HitlEvents",
     "PauseInfo",
     "PausedRunRecord",
     "PausedRunStore",
-    "agui_tool_to_external_function",
-    "emit_awaiting_tool_result",
-    "emit_messages_snapshot",
-    "make_tools_factory",
 ]

digitalkin-1.0.0.dev12/src/digitalkin/community/agno/agui_tools.py

@@ -0,0 +1,90 @@
+"""AG-UI frontend tools → Agno external Functions.
+
+The AG-UI protocol lets the client declare its own tools in
+``RunAgentInput.tools``, meant to be executed on the frontend rather than
+by the agent process. :class:`AguiTools` exposes them to an Agno agent as
+:class:`~agno.tools.function.Function` objects marked
+``external_execution=True``: when the LLM "calls" one, Agno pauses the run
+instead of executing an entrypoint, letting the caller stream the tool-call
+events to the front and resume later.
+
+See ``examples/`` and the :class:`~digitalkin.community.agno.AgnoHitlRunner`
+docstring for end-to-end usage.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from ag_ui.core.types import Tool as AgUiTool
+    from agno.run.base import RunContext
+    from agno.tools.function import Function
+
+
+class AguiTools:
+    """Convert AG-UI frontend tool declarations into Agno external Functions."""
+
+    @staticmethod
+    def _unreachable_entrypoint(**_: Any) -> None:
+        """Placeholder — never invoked because ``external_execution=True`` pauses the run."""
+
+    @staticmethod
+    def agui_tool_to_external_function(tool: AgUiTool) -> Function:
+        """Wrap an AG-UI tool definition as an Agno external ``Function``.
+
+        The resulting :class:`Function` carries the AG-UI schema as-is and is
+        marked ``external_execution=True`` so Agno emits the tool-call events
+        but skips the entrypoint and pauses the run when the LLM invokes it.
+
+        Args:
+            tool: An :class:`ag_ui.core.types.Tool` from ``RunAgentInput.tools``.
+
+        Returns:
+            An :class:`agno.tools.function.Function` ready to plug into an agent.
+        """
+        from agno.tools.function import Function  # pyright: ignore[reportMissingImports]
+
+        parameters = tool.parameters or {"type": "object", "properties": {}, "required": []}
+        return Function(
+            name=tool.name,
+            description=tool.description,
+            parameters=parameters,
+            entrypoint=AguiTools._unreachable_entrypoint,
+            external_execution=True,
+            skip_entrypoint_processing=True,
+        )
+
+    @staticmethod
+    def make_tools_factory(
+        base_tools: list[Any],
+        dependency_key: str = "agui_tools",
+    ) -> Callable[[RunContext], list[Any]]:
+        """Build an Agno ``tools`` factory merging base tools with per-run AG-UI tools.
+
+        The returned callable is the value passed to ``Agent(tools=...)``. On
+        every run Agno resolves it with the current ``RunContext``; the factory
+        reads ``run_context.dependencies[dependency_key]`` (the per-run AG-UI
+        tool list), converts them to external Functions, and concatenates them
+        with ``base_tools``.
+
+        Args:
+            base_tools: Toolkits / Functions always available, passed through.
+            dependency_key: Key in ``run_context.dependencies`` for the per-run
+                AG-UI tool list. Defaults to ``"agui_tools"``.
+
+        Returns:
+            A callable for ``Agent(tools=...)``. Set ``cache_callables=False``
+            so it is re-invoked every run.
+        """
+
+        def factory(run_context: RunContext | None = None) -> list[Any]:
+            if run_context is None:
+                return list(base_tools)
+            deps = getattr(run_context, "dependencies", None) or {}
+            agui_tools: list[AgUiTool] = deps.get(dependency_key) or []
+            return [*base_tools, *[AguiTools.agui_tool_to_external_function(t) for t in agui_tools]]
+
+        return factory

{digitalkin-1.0.0.dev10 → digitalkin-1.0.0.dev12}/src/digitalkin/community/agno/hitl.py

@@ -124,169 +124,173 @@ class PausedRunStore:
         await self._storage.remove(collection=self.COLLECTION, record_id=thread_id)
 
 
-def _agno_messages_to_agui(agno_messages: list[Any]) -> list[AgUiMessage]:
-    """Convert Agno messages into AG-UI messages.
-
-    Drops system/developer/reasoning; reshapes assistant ``tool_calls``
-    into AG-UI :class:`~ag_ui.core.types.ToolCall` objects.
-
-    Args:
-        agno_messages: Value of ``RunOutput.messages`` at pause time.
-
-    Returns:
-        AG-UI :class:`~ag_ui.core.types.Message` instances.
-    """
-    from ag_ui.core.types import (
-        AssistantMessage as AgUiAssistantMessage,
-    )
-    from ag_ui.core.types import (
-        FunctionCall as AgUiFunctionCall,
-    )
-    from ag_ui.core.types import (
-        ToolCall as AgUiToolCall,
-    )
-    from ag_ui.core.types import (
-        ToolMessage as AgUiToolMessage,
-    )
-    from ag_ui.core.types import (
-        UserMessage as AgUiUserMessage,
-    )
-
-    result: list[AgUiMessage] = []
-    for msg in agno_messages or []:
-        role = getattr(msg, "role", None)
-        msg_id = getattr(msg, "id", None) or ""
-        content = getattr(msg, "content", None)
-        if isinstance(content, list):
-            content = " ".join(str(part) for part in content if part is not None)
-
-        if role == "user":
-            result.append(AgUiUserMessage(id=msg_id, role="user", content=content or ""))
-        elif role == "assistant":
-            raw_calls = getattr(msg, "tool_calls", None) or []
-            agui_tool_calls: list[AgUiToolCall] = []
-            for tc in raw_calls:
-                tc_id = tc.get("id") if isinstance(tc, dict) else getattr(tc, "id", None)
-                func = tc.get("function") if isinstance(tc, dict) else getattr(tc, "function", None)
-                if not tc_id or func is None:
-                    continue
-                func_name = func.get("name") if isinstance(func, dict) else getattr(func, "name", None)
-                func_args = func.get("arguments") if isinstance(func, dict) else getattr(func, "arguments", None)
-                if not isinstance(func_args, str):
-                    func_args = json.dumps(func_args) if func_args is not None else "{}"
-                agui_tool_calls.append(
-                    AgUiToolCall(
-                        id=tc_id,
-                        type="function",
-                        function=AgUiFunctionCall(name=func_name or "", arguments=func_args),
+class HitlEvents:
+    """AG-UI message conversion and event emission for the HITL flow."""
+
+    @staticmethod
+    def agno_messages_to_agui(agno_messages: list[Any]) -> list[AgUiMessage]:
+        """Convert Agno messages into AG-UI messages.
+
+        Drops system/developer/reasoning; reshapes assistant ``tool_calls``
+        into AG-UI :class:`~ag_ui.core.types.ToolCall` objects.
+
+        Args:
+            agno_messages: Value of ``RunOutput.messages`` at pause time.
+
+        Returns:
+            AG-UI :class:`~ag_ui.core.types.Message` instances.
+        """
+        from ag_ui.core.types import (
+            AssistantMessage as AgUiAssistantMessage,
+        )
+        from ag_ui.core.types import (
+            FunctionCall as AgUiFunctionCall,
+        )
+        from ag_ui.core.types import (
+            ToolCall as AgUiToolCall,
+        )
+        from ag_ui.core.types import (
+            ToolMessage as AgUiToolMessage,
+        )
+        from ag_ui.core.types import (
+            UserMessage as AgUiUserMessage,
+        )
+
+        result: list[AgUiMessage] = []
+        for msg in agno_messages or []:
+            role = getattr(msg, "role", None)
+            msg_id = getattr(msg, "id", None) or ""
+            content = getattr(msg, "content", None)
+            if isinstance(content, list):
+                content = " ".join(str(part) for part in content if part is not None)
+
+            if role == "user":
+                result.append(AgUiUserMessage(id=msg_id, role="user", content=content or ""))
+            elif role == "assistant":
+                raw_calls = getattr(msg, "tool_calls", None) or []
+                agui_tool_calls: list[AgUiToolCall] = []
+                for tc in raw_calls:
+                    tc_id = tc.get("id") if isinstance(tc, dict) else getattr(tc, "id", None)
+                    func = tc.get("function") if isinstance(tc, dict) else getattr(tc, "function", None)
+                    if not tc_id or func is None:
+                        continue
+                    func_name = func.get("name") if isinstance(func, dict) else getattr(func, "name", None)
+                    func_args = func.get("arguments") if isinstance(func, dict) else getattr(func, "arguments", None)
+                    if not isinstance(func_args, str):
+                        func_args = json.dumps(func_args) if func_args is not None else "{}"
+                    agui_tool_calls.append(
+                        AgUiToolCall(
+                            id=tc_id,
+                            type="function",
+                            function=AgUiFunctionCall(name=func_name or "", arguments=func_args),
+                        )
+                    )
+                result.append(
+                    AgUiAssistantMessage(
+                        id=msg_id,
+                        role="assistant",
+                        content=content if isinstance(content, str) else None,
+                        tool_calls=agui_tool_calls or None,
                     )
                 )
-            result.append(
-                AgUiAssistantMessage(
-                    id=msg_id,
-                    role="assistant",
-                    content=content if isinstance(content, str) else None,
-                    tool_calls=agui_tool_calls or None,
+            elif role == "tool":
+                tool_call_id = getattr(msg, "tool_call_id", None)
+                if not tool_call_id:
+                    continue
+                result.append(
+                    AgUiToolMessage(
+                        id=msg_id,
+                        role="tool",
+                        tool_call_id=tool_call_id,
+                        content=content if isinstance(content, str) else "",
+                    )
                 )
+        return result
+
+    @staticmethod
+    async def emit_messages_snapshot(
+        context: ModuleContext,
+        messages: list[AgUiMessage],
+    ) -> None:
+        """Emit an AG-UI ``MessagesSnapshot`` event.
+
+        Typically called just before :func:`emit_awaiting_tool_result` on a
+        paused run so the front has an authoritative view of the conversation
+        (including the assistant message carrying the frontend ``tool_calls``,
+        which cannot be reconstructed from the streamed tool-call events alone).
+
+        Args:
+            context: Current module context.
+            messages: List of AG-UI messages, typically produced by
+                :func:`agno_messages_to_agui` from ``RunOutput.messages``.
+        """
+        if not messages:
+            return
+
+        from ag_ui.core.events import MessagesSnapshotEvent as AgUiMessagesSnapshotEvent
+
+        from digitalkin.models.module.ag_ui import (
+            AgUiMessagesSnapshotOutput,
+            AgUiOutput,
+        )
+
+        output = AgUiOutput(
+            root=AgUiMessagesSnapshotOutput(
+                event=AgUiMessagesSnapshotEvent(messages=messages),
             )
-        elif role == "tool":
-            tool_call_id = getattr(msg, "tool_call_id", None)
-            if not tool_call_id:
-                continue
-            result.append(
-                AgUiToolMessage(
-                    id=msg_id,
-                    role="tool",
-                    tool_call_id=tool_call_id,
-                    content=content if isinstance(content, str) else "",
+        )
+        await context.callbacks.send_message(output)
+        logger.info("emit_messages_snapshot: sent %d message(s)", len(messages))
+
+    @staticmethod
+    async def emit_awaiting_tool_result(
+        context: ModuleContext,
+        *,
+        thread_id: str,
+        run_id: str,
+        pending_tool_call_ids: list[str],
+    ) -> None:
+        """Emit an AG-UI ``RunFinished`` with ``status="awaiting_tool_result"``.
+
+        This is the protocol signal telling the front "the run paused on a
+        client-side tool; execute it and reply with a ``ToolMessage``". It
+        goes out via ``context.callbacks.send_message`` (bypassing the
+        standard :class:`~digitalkin.mixins.agui_mixin.AgUiMixin` event
+        mapping, which has no notion of an "awaiting" status).
+
+        Args:
+            context: Current module context.
+            thread_id: AG-UI thread identifier.
+            run_id: Run identifier to echo back in the finished event.
+            pending_tool_call_ids: The ``tool_call_id`` values the front must
+                execute and resolve — echoed in ``result.pending_tool_call_ids``
+                so the front can match them.
+        """
+        from ag_ui.core.events import RunFinishedEvent as AgUiRunFinishedEvent
+
+        from digitalkin.models.module.ag_ui import (
+            AgUiOutput,
+            AgUiRunFinishedOutput,
+        )
+
+        output = AgUiOutput(
+            root=AgUiRunFinishedOutput(
+                event=AgUiRunFinishedEvent(
+                    thread_id=thread_id,
+                    run_id=run_id,
+                    result={
+                        "status": _AWAITING_STATUS,
+                        "pending_tool_call_ids": pending_tool_call_ids,
+                    },
                 )
             )
-    return result
-
-
-async def emit_messages_snapshot(
-    context: ModuleContext,
-    messages: list[AgUiMessage],
-) -> None:
-    """Emit an AG-UI ``MessagesSnapshot`` event.
-
-    Typically called just before :func:`emit_awaiting_tool_result` on a
-    paused run so the front has an authoritative view of the conversation
-    (including the assistant message carrying the frontend ``tool_calls``,
-    which cannot be reconstructed from the streamed tool-call events alone).
-
-    Args:
-        context: Current module context.
-        messages: List of AG-UI messages, typically produced by
-            :func:`_agno_messages_to_agui` from ``RunOutput.messages``.
-    """
-    if not messages:
-        return
-
-    from ag_ui.core.events import MessagesSnapshotEvent as AgUiMessagesSnapshotEvent
-
-    from digitalkin.models.module.ag_ui import (
-        AgUiMessagesSnapshotOutput,
-        AgUiOutput,
-    )
-
-    output = AgUiOutput(
-        root=AgUiMessagesSnapshotOutput(
-            event=AgUiMessagesSnapshotEvent(messages=messages),
         )
-    )
-    await context.callbacks.send_message(output)
-    logger.info("emit_messages_snapshot: sent %d message(s)", len(messages))
-
-
-async def emit_awaiting_tool_result(
-    context: ModuleContext,
-    *,
-    thread_id: str,
-    run_id: str,
-    pending_tool_call_ids: list[str],
-) -> None:
-    """Emit an AG-UI ``RunFinished`` with ``status="awaiting_tool_result"``.
-
-    This is the protocol signal telling the front "the run paused on a
-    client-side tool; execute it and reply with a ``ToolMessage``". It
-    goes out via ``context.callbacks.send_message`` (bypassing the
-    standard :class:`~digitalkin.mixins.agui_mixin.AgUiMixin` event
-    mapping, which has no notion of an "awaiting" status).
-
-    Args:
-        context: Current module context.
-        thread_id: AG-UI thread identifier.
-        run_id: Run identifier to echo back in the finished event.
-        pending_tool_call_ids: The ``tool_call_id`` values the front must
-            execute and resolve — echoed in ``result.pending_tool_call_ids``
-            so the front can match them.
-    """
-    from ag_ui.core.events import RunFinishedEvent as AgUiRunFinishedEvent
-
-    from digitalkin.models.module.ag_ui import (
-        AgUiOutput,
-        AgUiRunFinishedOutput,
-    )
-
-    output = AgUiOutput(
-        root=AgUiRunFinishedOutput(
-            event=AgUiRunFinishedEvent(
-                thread_id=thread_id,
-                run_id=run_id,
-                result={
-                    "status": _AWAITING_STATUS,
-                    "pending_tool_call_ids": pending_tool_call_ids,
-                },
-            )
+        await context.callbacks.send_message(output)
+        logger.info(
+            "emit_awaiting_tool_result: thread_id=%s pending=%s",
+            thread_id,
+            pending_tool_call_ids,
         )
-    )
-    await context.callbacks.send_message(output)
-    logger.info(
-        "emit_awaiting_tool_result: thread_id=%s pending=%s",
-        thread_id,
-        pending_tool_call_ids,
-    )
 
 
 class AgnoHitlRunner:
@@ -559,7 +563,7 @@ class AgnoHitlRunner:
         resumed, pause_info = await self.try_resume(input_data=input_data, send=send)
         if resumed:
             if pause_info is not None and context is not None:
-                await emit_awaiting_tool_result(
+                await HitlEvents.emit_awaiting_tool_result(
                     context,
                     thread_id=pause_info.thread_id,
                     run_id=pause_info.run_id,
@@ -586,7 +590,7 @@ class AgnoHitlRunner:
             images=images,
         )
         if pause_info is not None and context is not None:
-            await emit_awaiting_tool_result(
+            await HitlEvents.emit_awaiting_tool_result(
                 context,
                 thread_id=pause_info.thread_id,
                 run_id=pause_info.run_id,
@@ -625,7 +629,7 @@ class AgnoHitlRunner:
         if adapter.is_paused and final_run_output is not None and getattr(final_run_output, "is_paused", False):
             pause_info = await self._store.save(run_output=final_run_output, thread_id=thread_id)
             # Attach AG-UI-shaped messages so the front can materialise the tool_call.
-            pause_info.new_messages = _agno_messages_to_agui(final_run_output.messages or [])
+            pause_info.new_messages = HitlEvents.agno_messages_to_agui(final_run_output.messages or [])
             return pause_info
 
         return None

{digitalkin-1.0.0.dev10 → digitalkin-1.0.0.dev12}/src/digitalkin/core/job_manager/base_job_manager.py

@@ -202,3 +202,37 @@ class BaseJobManager(abc.ABC, Generic[InputModelT, OutputModelT, SetupModelT]):
         Returns:
             dict[str, dict[str, Any]]: A dictionary containing information about all modules and their statuses.
         """
+
+    @abc.abstractmethod
+    async def preload_instance(
+        self,
+        setup_data: SetupModelT,
+        mission_id: str,
+        setup_id: str,
+        setup_version_id: str,
+        request_metadata: dict[str, str] | None = None,
+        job_id: str | None = None,
+        tool_cache: Any = None,
+        callback: Callable | None = None,
+    ) -> tuple[Any, str, Callable]:
+        """Build a module instance and run its idempotent ``prepare()``.
+
+        Returns:
+            Tuple of (prepared module instance, job_id, output callback).
+        """
+
+    @abc.abstractmethod
+    async def run_instance(
+        self,
+        module: Any,
+        job_id: str,
+        mission_id: str,
+        input_data: InputModelT,
+        setup_data: SetupModelT,
+        callback: Callable,
+    ) -> str:
+        """Run a pre-prepared module instance (from ``preload_instance``) with input.
+
+        Returns:
+            The job_id of the scheduled run.
+        """

{digitalkin-1.0.0.dev10 → digitalkin-1.0.0.dev12}/src/digitalkin/core/job_manager/single_job_manager.py

@@ -273,6 +273,8 @@ class SingleJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
         timer.mark("redis_task_manager")
 
         if callback is None:
+            if self._stream_writers is None:  # safety when __init__ was bypassed (e.g. object.__new__)
+                self._stream_writers = {}
             self._stream_writers[job_id] = RedisStreamWriter(job_id, self._redis_client)
             callback = await self.job_specific_callback(self.add_to_queue, job_id)
             timer.mark("default_callback")

{digitalkin-1.0.0.dev10 → digitalkin-1.0.0.dev12}/src/digitalkin/core/task_manager/module_runner.py

@@ -14,9 +14,7 @@ from digitalkin.core.profiling.step_timer import StepTimer
 from digitalkin.core.profiling.task_profiler import TaskProfiler
 from digitalkin.logger import logger
 from digitalkin.models.grpc_servers.stream_error_codes import StreamErrorCode
-from digitalkin.models.settings.profiling import ProfilerMode, ProfilingSettings
-
-_PROFILING = ProfilingSettings()
+from digitalkin.models.settings.profiling import ProfilerMode, get_profiling_settings
 
 if TYPE_CHECKING:
     from collections.abc import Awaitable, Callable
@@ -65,12 +63,13 @@ class ModuleRunner:
         timer = StepTimer()
 
         # Construct profiler outside try so finally can always stop it.
+        profiling = get_profiling_settings()
         profiler_mode = (
-            ProfilerMode(_PROFILING.profiler)
-            if _PROFILING.profiler in {p.value for p in ProfilerMode}
+            ProfilerMode(profiling.profiler)
+            if profiling.profiler in {p.value for p in ProfilerMode}
             else ProfilerMode.NONE
         )
-        profiler = TaskProfiler(task_id=task_id, mode=profiler_mode, output_dir=_PROFILING.profile_output_dir)
+        profiler = TaskProfiler(task_id=task_id, mode=profiler_mode, output_dir=profiling.profile_output_dir)
 
         try:
             timer.mark("entry")

{digitalkin-1.0.0.dev10 → digitalkin-1.0.0.dev12}/src/digitalkin/core/task_manager/redis/instrumented.py

@@ -15,8 +15,6 @@ if TYPE_CHECKING:
 
 from digitalkin.logger import logger
 
-_MIN_KEY_PARTS_FOR_REDACTION = 2
-
 
 class InstrumentedRedisClient:  # noqa: PLR0904
     """Observability wrapper around RedisClient.
@@ -52,8 +50,8 @@ class InstrumentedRedisClient:  # noqa: PLR0904
         parts = key.split(":")
         if len(parts) <= 1:
             return key
-        if len(parts) > _MIN_KEY_PARTS_FOR_REDACTION:
-            return ":".join(parts[0:1] + ["*"] * (len(parts) - _MIN_KEY_PARTS_FOR_REDACTION) + parts[-1:])
+        if len(parts) > 2:  # noqa: PLR2004
+            return ":".join(parts[0:1] + ["*"] * (len(parts) - 2) + parts[-1:])
         return f"{parts[0]}:*"
 
     async def _execute(self, command: str, key: str, coro: Any) -> Any: