Fix/Organize Types for ModuleCoordiator/DimosCluster/WorkerManager/RPCClient/ModuleProxy

README.md

@@ -189,7 +189,7 @@ if __name__ == "__main__":
 ### Blueprints
 
 Blueprints are how robots are constructed on Dimensional; instructions for how to construct and wire modules. You compose them with
-`autoconnect(...)`, which connects streams by `(name, type)` and returns a `ModuleBlueprintSet`.
+`autoconnect(...)`, which connects streams by `(name, type)` and returns a `Blueprint`.
 
 Blueprints can be composed, remapped, and have transports overridden if `autoconnect()` fails due to conflicting variable names or `In[]` and `Out[]` message types.
 

dimos/agents/agent.py

@@ -29,7 +29,7 @@
 
 from dimos.agents.llm_init import build_llm, build_system_message
 from dimos.agents.spec import AgentSpec, Model, Provider
-from dimos.core import DimosCluster, rpc
+from dimos.core import DaskDeployer, rpc
 from dimos.protocol.skill.coordinator import SkillCoordinator, SkillState, SkillStateDict
 from dimos.protocol.skill.skill import SkillContainer
 from dimos.protocol.skill.type import Output
@@ -369,7 +369,7 @@ def stop(self) -> None:
 
 
 def deploy(
-    dimos: DimosCluster,
+    dimos: DaskDeployer,
     system_prompt: str = "You are a helpful assistant for controlling a Unitree Go2 robot.",
     model: Model = Model.GPT_4O,
     provider: Provider = Provider.OPENAI,  # type: ignore[attr-defined]

dimos/core/__init__.py

@@ -3,12 +3,14 @@
 import multiprocessing as mp
 import signal
 import time
+from typing import TYPE_CHECKING, cast
 
 from dask.distributed import Client, LocalCluster
 from rich.console import Console
 
 import dimos.core.colors as colors
 from dimos.core.core import rpc
+from dimos.core.deployer_protocol import DeployerProtocol, ModuleProxy
 from dimos.core.module import Module, ModuleBase, ModuleConfig, ModuleConfigT
 from dimos.core.rpc_client import RPCClient
 from dimos.core.stream import In, Out, RemoteIn, RemoteOut, Transport
@@ -31,7 +33,7 @@
     "LCMRPC",
     "LCMTF",
     "TF",
-    "DimosCluster",
+    "DaskDeployer",
     "In",
     "LCMTransport",
     "Module",
@@ -86,15 +88,15 @@ def teardown(self, worker) -> None:  # type: ignore[no-untyped-def]
 def patch_actor(actor, cls) -> None: ...  # type: ignore[no-untyped-def]
 
 
-DimosCluster = Client
+class DaskDeployer(Client, DeployerProtocol): ...
 
 
-def patchdask(dask_client: Client, local_cluster: LocalCluster) -> DimosCluster:
+def patchdask(dask_client: Client, local_cluster: LocalCluster) -> DaskDeployer:
     def deploy(  # type: ignore[no-untyped-def]
-        actor_class,
+        actor_class: type[Module],
         *args,
         **kwargs,
-    ):
+    ) -> ModuleProxy:
         logger.info("Deploying module.", module=actor_class.__name__)
         actor = dask_client.submit(  # type: ignore[no-untyped-call]
             actor_class,
@@ -109,7 +111,7 @@ def deploy(  # type: ignore[no-untyped-def]
         # Register actor deployment in shared memory
         ActorRegistry.update(str(actor), str(worker))
 
-        return RPCClient(actor, actor_class)
+        return cast("ModuleProxy", RPCClient(actor, actor_class))
 
     def check_worker_memory() -> None:
         """Check memory usage of all workers."""
@@ -225,18 +227,18 @@ def close_all() -> None:
     dask_client.check_worker_memory = check_worker_memory  # type: ignore[attr-defined]
     dask_client.stop = lambda: dask_client.close()  # type: ignore[attr-defined, no-untyped-call]
     dask_client.close_all = close_all  # type: ignore[attr-defined]
-    return dask_client
+    return dask_client  # type: ignore[return-value]
 
 
-def start(n: int | None = None, memory_limit: str = "auto") -> DimosCluster:
+def start(n: int | None = None, memory_limit: str = "auto") -> DaskDeployer:
     """Start a Dask LocalCluster with specified workers and memory limits.
 
     Args:
         n: Number of workers (defaults to CPU count)
         memory_limit: Memory limit per worker (e.g., '4GB', '2GiB', or 'auto' for Dask's default)
 
     Returns:
-        DimosCluster: A patched Dask client with deploy(), check_worker_memory(), stop(), and close_all() methods
+        DaskDeployer: A patched Dask client with deploy(), check_worker_memory(), stop(), and close_all() methods
     """
 
     console = Console()