refactor: replace monkey-patching with dependency injection in train() — closes #64

clawloop/cli.py

@@ -18,32 +18,9 @@
 from pathlib import Path
 from typing import Any
 
-from clawloop.train import LLMClientConfig
-
 log = logging.getLogger("clawloop")
 
 
-class _DryRunLLMClientConfig(LLMClientConfig):
-    """Private subclass that tags an LLM client config with its dry-run role.
-
-    Used only by ``cmd_run --dry-run``: ``_install_dry_run_clients`` swaps
-    each entry in ``config.llm_clients`` for an instance of this class, and
-    the patched ``_make_llm_client`` picks the right mock via ``isinstance``.
-
-    Why a subclass instead of a field on ``LLMClientConfig``:
-      - The public ``LLMClientConfig`` schema stays free of testing
-        vocabulary — no ``dry_run_role: null`` in JSON dumps or generated
-        JSON Schema.
-      - ``model`` is left untouched, so downstream code that reads it
-        verbatim (e.g. ``_build_entropic`` propagating ``tc.model`` into
-        ``entropic_cfg``) is unaffected.
-      - Pydantic ``model_copy()`` preserves the runtime class, so the tag
-        survives copies just like a field would.
-    """
-
-    dry_run_role: str
-
-
 _EVAL_DISABLED_MSG = (
     "`clawloop eval` is disabled. Use one of:\n"
     "  - Real benchmark:  uv run clawloop run examples/configs/math_harness.json\n"
@@ -180,69 +157,65 @@ def cmd_run(args: argparse.Namespace) -> None:
     )
 
     if args.dry_run:
-        _install_dry_run_clients(config)
+        factory, builders = _build_dry_run_overrides(config)
+        train(config, make_llm_client=factory, env_builders=builders)
+    else:
+        train(config)
 
-    train(config)
 
+def _build_dry_run_overrides(config: Any) -> tuple[Any, dict[str, Any]]:
+    """Build the `(make_llm_client, env_builders)` overrides for `--dry-run`.
 
-def _install_dry_run_clients(config: Any) -> None:
-    """Wire `--dry-run`: guarantee no real LLM calls regardless of env_type.
+    Pure function: returns the overrides without mutating any module-level
+    state. ``train()`` then accepts them as keyword arguments and uses them
+    in place of the global defaults. This replaces the earlier
+    monkey-patching approach (see PR #62, PR #65 history).
 
     Two parts:
-      1. Replace each ``LLMClientConfig`` in ``config.llm_clients`` with a
-         ``_DryRunLLMClientConfig`` instance carrying the role, and patch
-         ``clawloop.train._make_llm_client`` to switch on ``isinstance``.
-         The private subclass keeps the public schema clean and survives
-         ``model_copy()``.
-      2. For envs whose adapter bypasses ``_make_llm_client`` (per
-         ``train.ENVS_USING_MAKE_LLM_CLIENT``), swap the registered builder
-         with a stub that returns a no-I/O ``_StubAdapter``.
+      1. A factory closure that maps each ``LLMClientConfig`` to the right
+         mock client based on its role in ``config.llm_clients``. The
+         closure looks up by object identity, which is safe here because
+         every builder receives ``config.llm_clients`` directly and indexes
+         by role name — no copies between cfg creation and factory call.
+      2. For envs whose builder bypasses ``make_llm_client`` (per
+         ``train.ENVS_USING_MAKE_LLM_CLIENT``), an override entry whose
+         builder returns a no-I/O ``_StubAdapter``.
     """
     import clawloop.train as _train
     from clawloop.demo_math import MockTaskClient, _build_mock_reflector_responses
     from clawloop.llm import MockLLMClient
 
-    # Part 1: replace each cfg with a subclass instance carrying the role,
-    # then route _make_llm_client through a mock factory that reads the role.
-    for role, cfg in list(config.llm_clients.items()):
-        config.llm_clients[role] = _DryRunLLMClientConfig(**cfg.model_dump(), dry_run_role=role)
-
-    original_make = _train._make_llm_client
-
-    def _mock_make(cfg):
-        if isinstance(cfg, _DryRunLLMClientConfig):
-            role = cfg.dry_run_role
-            if role == "reflector":
-                return MockLLMClient(responses=_build_mock_reflector_responses())
-            if role == "task":
-                return MockTaskClient()
-            return MockLLMClient(responses=["[]"])
-        return original_make(cfg)
-
-    _train._make_llm_client = _mock_make
-
-    # Part 2: for env_types that bypass _make_llm_client, replace the
-    # registered builder with one that returns a stub adapter. Without this,
-    # --dry-run on (e.g.) taubench / entropic / openclaw would still hit
-    # real endpoints.
+    role_by_id = {id(cfg): role for role, cfg in config.llm_clients.items()}
+
+    def factory(cfg: Any) -> Any:
+        role = role_by_id.get(id(cfg))
+        if role == "reflector":
+            return MockLLMClient(responses=_build_mock_reflector_responses())
+        if role == "task":
+            return MockTaskClient()
+        return MockLLMClient(responses=["[]"])
+
     env_type = config.env_type
     uses_make_llm_client = env_type in _train.ENVS_USING_MAKE_LLM_CLIENT
-    if not uses_make_llm_client and env_type in _train.ENV_BUILDERS:
+    if uses_make_llm_client or env_type not in _train.ENV_BUILDERS:
+        builders: dict[str, Any] = _train.ENV_BUILDERS
+    else:
         # Floor at 1 to keep the task list non-empty even if a config sets
         # episodes_per_iter to 0; the learning loop samples from it.
         n_tasks = max(1, config.episodes_per_iter)
         stub_tasks = [f"dry_run_{env_type}_{i}" for i in range(n_tasks)]
 
-        def _stub_builder(_cfg: Any, _clients: Any) -> tuple[Any, list[str]]:
+        def stub_builder(_cfg: Any, _clients: Any, _factory: Any) -> tuple[Any, list[str]]:
             return _StubAdapter(env_type), list(stub_tasks)
 
-        _train.ENV_BUILDERS[env_type] = _stub_builder
+        builders = {**_train.ENV_BUILDERS, env_type: stub_builder}
 
     log.info(
         "dry-run: LLM clients mocked; env=%r %s",
         env_type,
-        "uses _make_llm_client" if uses_make_llm_client else "stubbed",
+        "uses make_llm_client" if uses_make_llm_client else "stubbed",
     )
+    return factory, builders
 
 
 class _StubAdapter:

clawloop/train.py

@@ -5,17 +5,24 @@
   harness_learning — prompt/playbook optimization via reflector LLM (no GPU)
 
 Environments are pluggable via ENV_BUILDERS registry. Each builder is a
-function (config, llm_clients) -> (adapter, tasks) that constructs an
-AdapterLike and a task list for the learning loop.
+function ``(config, llm_clients, make_llm_client) -> (adapter, tasks)``
+that constructs an AdapterLike and a task list for the learning loop.
+
+``train()`` accepts ``make_llm_client`` and ``env_builders`` as keyword-only
+parameters so callers (CLI dry-run, tests) can inject mocks without
+mutating module-level state.
 """
 
 from __future__ import annotations
 
+from collections.abc import Callable
 from pathlib import Path
 from typing import Any, Literal
 
 from pydantic import BaseModel, SecretStr
 
+LLMClientFactory = Callable[["LLMClientConfig"], Any]
+
 # ---------------------------------------------------------------------------
 # Config
 # ---------------------------------------------------------------------------
@@ -98,7 +105,11 @@ def _make_llm_client(cfg: LLMClientConfig):
 # ---------------------------------------------------------------------------
 
 
-def _build_harbor(config: TrainConfig, llm_clients: dict[str, LLMClientConfig]):
+def _build_harbor(
+    config: TrainConfig,
+    llm_clients: dict[str, LLMClientConfig],
+    make_llm_client: LLMClientFactory,
+) -> tuple[Any, list[str]]:
     from clawloop.environments.harbor import HarborAdapter, HarborTaskEnvironment
 
     envs = [
@@ -112,16 +123,24 @@ def _build_harbor(config: TrainConfig, llm_clients: dict[str, LLMClientConfig]):
     return HarborAdapter(envs), [env.task_id for env in envs]
 
 
-def _build_math(config: TrainConfig, llm_clients: dict[str, LLMClientConfig]):
+def _build_math(
+    config: TrainConfig,
+    llm_clients: dict[str, LLMClientConfig],
+    make_llm_client: LLMClientFactory,
+) -> tuple[Any, list[str]]:
     from clawloop.environments.math import MathAdapter, MathEnvironment
 
-    task_client = _make_llm_client(llm_clients["task"])
+    task_client = make_llm_client(llm_clients["task"])
     math_env = MathEnvironment()
     tasks = math_env.get_tasks()
     return MathAdapter(env=math_env, client=task_client), [s.question for s in tasks]
 
 
-def _build_entropic(config: TrainConfig, llm_clients: dict[str, LLMClientConfig]):
+def _build_entropic(
+    config: TrainConfig,
+    llm_clients: dict[str, LLMClientConfig],
+    make_llm_client: LLMClientFactory,
+) -> tuple[Any, list[str]]:
     from clawloop.environments.entropic import EntropicAdapter
 
     entropic_cfg = dict(config.env_config or {})
@@ -141,7 +160,11 @@ def _build_entropic(config: TrainConfig, llm_clients: dict[str, LLMClientConfig]
     return adapter, [f"base_{i}" for i in range(n_tasks)]
 
 
-def _build_openclaw(config: TrainConfig, llm_clients: dict[str, LLMClientConfig]) -> tuple:
+def _build_openclaw(
+    config: TrainConfig,
+    llm_clients: dict[str, LLMClientConfig],
+    make_llm_client: LLMClientFactory,
+) -> tuple[Any, list[str]]:
     from clawloop.environments.openclaw import OpenClawAdapter
 
     openclaw_cfg = dict(config.env_config or {})
@@ -170,7 +193,9 @@ def _build_openclaw(config: TrainConfig, llm_clients: dict[str, LLMClientConfig]
 
 
 def _build_taubench(
-    config: TrainConfig, llm_clients: dict[str, LLMClientConfig]
+    config: TrainConfig,
+    llm_clients: dict[str, LLMClientConfig],
+    make_llm_client: LLMClientFactory,
 ) -> tuple[Any, list[str]]:
     from clawloop.environments.taubench import TauBenchAdapter
 
@@ -189,7 +214,11 @@ def _build_taubench(
 # ---------------------------------------------------------------------------
 
 
-def _build_openspiel(config: "TrainConfig", llm_clients: dict[str, "LLMClientConfig"]):
+def _build_openspiel(
+    config: "TrainConfig",
+    llm_clients: dict[str, "LLMClientConfig"],
+    make_llm_client: LLMClientFactory,
+) -> tuple[Any, list[str]]:
     """Build a ClawLoop adapter over one or more OpenSpiel games.
 
     Two config shapes are accepted:
@@ -425,19 +454,32 @@ def _positive_int(key: str) -> None:
 # ---------------------------------------------------------------------------
 
 
-def train(config: TrainConfig):
+def train(
+    config: TrainConfig,
+    *,
+    make_llm_client: LLMClientFactory | None = None,
+    env_builders: dict[str, Any] | None = None,
+):
     """Unified training entry point.
 
     Builds layers, environment adapter, and runs the learning loop.
     Environment is selected via env_type and constructed by the matching
-    builder from ENV_BUILDERS.
+    builder from ``env_builders`` (defaults to the module-level
+    ``ENV_BUILDERS``).
+
+    Both ``make_llm_client`` and ``env_builders`` are keyword-only so
+    callers (CLI dry-run, tests) can inject mocks / stubs without
+    mutating module-level state. Defaults preserve current behavior.
     """
     from clawloop.core.intensity import AdaptiveIntensity
     from clawloop.core.loop import AgentState, learning_loop
     from clawloop.learning_layers.harness import Harness
     from clawloop.learning_layers.router import Router
     from clawloop.learning_layers.weights import Weights
 
+    factory: LLMClientFactory = make_llm_client or _make_llm_client
+    builders = env_builders if env_builders is not None else ENV_BUILDERS
+
     layers = validate_config(config)
 
     # 1. Harness — reflector wired via LocalEvolver when harness layer is active
@@ -450,7 +492,7 @@ def train(config: TrainConfig):
         from clawloop.core.reflector import Reflector
         from clawloop.harness_backends.local import LocalEvolver
 
-        reflector_client = _make_llm_client(config.llm_clients["reflector"])
+        reflector_client = factory(config.llm_clients["reflector"])
         evolver = LocalEvolver(reflector=Reflector(client=reflector_client))
 
     harness = Harness(system_prompts=prompts, evolver=evolver)
@@ -482,8 +524,8 @@ def train(config: TrainConfig):
     router = Router()
 
     # 3. Environment adapter — dispatched via registry
-    build_env = ENV_BUILDERS[config.env_type]
-    adapter, tasks = build_env(config, config.llm_clients)
+    build_env = builders[config.env_type]
+    adapter, tasks = build_env(config, config.llm_clients, factory)
 
     # Wire harness into adapter for proxy skill injection (openclaw)
     if hasattr(adapter, "set_harness"):