From 08932dba052027dec39741967b227f34a6cb1f46 Mon Sep 17 00:00:00 2001 From: Kavya Sree Kaitepalli Date: Tue, 11 Aug 2026 05:33:54 +0000 Subject: [PATCH] improve feedback loop --- src/microbots/auto_memory/__init__.py | 48 ++- src/microbots/auto_memory/__main__.py | 11 + src/microbots/auto_memory/callbacks.py | 89 ++++- src/microbots/auto_memory/cli.py | 227 +++++------ src/microbots/auto_memory/config.py | 118 +++--- src/microbots/auto_memory/loop.py | 193 ++++++++++ src/microbots/auto_memory/memory.py | 30 +- src/microbots/auto_memory/orchestrator.py | 186 +++++---- src/microbots/auto_memory/runners/__init__.py | 3 +- src/microbots/auto_memory/runners/base.py | 8 +- .../auto_memory/runners/writing_bot_runner.py | 4 +- .../auto_memory/training/__init__.py | 2 + .../auto_memory/training/orchestrator.py | 161 +++++--- src/microbots/auto_memory/training/runner.py | 13 +- src/microbots/auto_memory/workspace.py | 19 +- .../runners/test_writing_bot_runner.py | 11 +- test/auto_memory/test_cli.py | 332 ++++++++-------- test/auto_memory/test_config.py | 107 ++--- test/auto_memory/test_context.py | 23 ++ test/auto_memory/test_loop.py | 145 +++++++ test/auto_memory/test_memory.py | 68 ++++ test/auto_memory/test_orchestrator.py | 23 +- test/auto_memory/test_workspace.py | 44 +++ test/swe-bench-test/run_plain_baseline.py | 364 ++++++++++++++++++ 24 files changed, 1632 insertions(+), 597 deletions(-) create mode 100644 src/microbots/auto_memory/__main__.py create mode 100644 src/microbots/auto_memory/loop.py create mode 100644 test/auto_memory/test_loop.py create mode 100644 test/swe-bench-test/run_plain_baseline.py diff --git a/src/microbots/auto_memory/__init__.py b/src/microbots/auto_memory/__init__.py index ec338cbe..76e2ca99 100644 --- a/src/microbots/auto_memory/__init__.py +++ b/src/microbots/auto_memory/__init__.py @@ -1,13 +1,57 @@ """Iterative agent loop with memory feedback (auto_memory package).""" +from microbots.auto_memory.callbacks import CallbackResult, CallbackRunner from microbots.auto_memory.cli import run_from_yaml -from microbots.auto_memory.config import TaskConfig -from microbots.auto_memory.orchestrator import TrainingLoopOrchestrator +from microbots.auto_memory.config import DEFAULT_PROMPT_TEMPLATE, TaskConfig +from microbots.auto_memory.data_models import ( + CallbackSpec, + Feedback, + FinalStatus, + IterationStatus, + ReferenceInput, +) +from microbots.auto_memory.errors import ( + AgentError, + AutoMemoryError, + AutoMemoryTimeoutError, + CallbackError, + ConfigError, + MemoryStoreError, +) +from microbots.auto_memory.memory import MemoryStore +from microbots.auto_memory.orchestrator import ( + IterationRecord, + RunSummary, + TrainingLoopOrchestrator, +) +from microbots.auto_memory.runners import AgentResult, AgentRunner, IterationContext from microbots.auto_memory.runners.writing_bot_runner import WritingBotRunner +from microbots.auto_memory.workspace import WorkspaceManager __all__ = [ + "CallbackResult", + "CallbackRunner", + "CallbackSpec", + "DEFAULT_PROMPT_TEMPLATE", + "AgentResult", + "AgentRunner", + "IterationContext", + "IterationRecord", + "IterationStatus", + "FinalStatus", + "Feedback", + "ReferenceInput", + "RunSummary", + "MemoryStore", + "WorkspaceManager", "TrainingLoopOrchestrator", "TaskConfig", "WritingBotRunner", "run_from_yaml", + "AutoMemoryError", + "ConfigError", + "AgentError", + "CallbackError", + "AutoMemoryTimeoutError", + "MemoryStoreError", ] \ No newline at end of file diff --git a/src/microbots/auto_memory/__main__.py b/src/microbots/auto_memory/__main__.py new file mode 100644 index 00000000..fb5ed995 --- /dev/null +++ b/src/microbots/auto_memory/__main__.py @@ -0,0 +1,11 @@ +"""Enable ``python -m microbots.auto_memory``.""" + +from __future__ import annotations + +import sys + +from microbots.auto_memory.cli import main + + +if __name__ == "__main__": # pragma: no cover + sys.exit(main()) \ No newline at end of file diff --git a/src/microbots/auto_memory/callbacks.py b/src/microbots/auto_memory/callbacks.py index 5cfb070b..d5118636 100644 --- a/src/microbots/auto_memory/callbacks.py +++ b/src/microbots/auto_memory/callbacks.py @@ -102,7 +102,18 @@ def run_all( list[CallbackResult] One result per spec, in the same order as *specs*. """ - return [self._run_one(spec, logs_dir, candidate_path) for spec in specs] + logger.info( + "Running %d callback(s) against candidate at %s", len(specs), candidate_path + ) + results = [self._run_one(spec, logs_dir, candidate_path) for spec in specs] + passed = sum(1 for r in results if r.passed) + logger.info( + "Callback batch finished: %d/%d passed (%.1fs total)", + passed, + len(results), + sum(r.duration_s for r in results), + ) + return results # ------------------------------------------------------------------ # Internal @@ -140,16 +151,19 @@ def _run_one( timed_out: bool = False start = time.monotonic() + logger.info( + "Callback %r starting: cmd=%s (timeout=%ds, expected_rc=%d)", + spec.name, + _preview(spec.command), + spec.timeout_s, + spec.expected_return_code, + ) try: with ( stdout_path.open("w", encoding="utf-8") as out_fh, stderr_path.open("w", encoding="utf-8") as err_fh, ): - # Security note: spec.command is intentionally run with - # shell=True for developer convenience (supports pipes, - # redirects, etc.). This runner assumes configs are loaded - # from trusted local files only. Do NOT use with configs - # sourced from untrusted input. + # spec.command runs with shell=True; assumes configs come from trusted local files. proc = subprocess.run( spec.command, shell=True, @@ -173,6 +187,25 @@ def _run_one( duration_s = time.monotonic() - start passed = (not timed_out) and (return_code == spec.expected_return_code) + logger.info( + "Callback %r finished: rc=%d elapsed=%.1fs passed=%s stdout=%s stderr=%s", + spec.name, + return_code, + duration_s, + passed, + stdout_path, + stderr_path, + ) + if not passed and not timed_out: + tail = _tail_lines(stderr_path, 15) or _tail_lines(stdout_path, 15) + if tail: + logger.info( + "Callback %r output tail (%d lines):\n%s", + spec.name, + len(tail), + "\n".join(tail), + ) + return CallbackResult( spec=spec, return_code=return_code, @@ -182,3 +215,47 @@ def _run_one( timed_out=timed_out, duration_s=duration_s, ) + + +def _preview(cmd: str, limit: int = 160) -> str: + """Truncate long callback commands so a single log line stays readable. + + Parameters + ---------- + cmd : str + Callback command to normalize and preview. + limit : int, optional + Maximum preview length before truncation. + + Returns + ------- + str + Single-line command preview. + """ + single = " ".join(cmd.split()) + if len(single) <= limit: + return single + return single[:limit] + f"... [{len(single) - limit} more chars]" + + +def _tail_lines(path: Path, n: int) -> list[str]: + """Read the last non-empty lines from a text file. + + Parameters + ---------- + path : Path + Text file to read. + n : int + Maximum number of trailing non-empty lines to return. + + Returns + ------- + list[str] + Trailing non-empty lines, or an empty list on an I/O error. + """ + try: + text = path.read_text(encoding="utf-8", errors="replace") + except OSError: + return [] + lines = [ln for ln in text.splitlines() if ln.strip()] + return lines[-n:] diff --git a/src/microbots/auto_memory/cli.py b/src/microbots/auto_memory/cli.py index 32c82043..a85f3b33 100644 --- a/src/microbots/auto_memory/cli.py +++ b/src/microbots/auto_memory/cli.py @@ -2,123 +2,33 @@ from __future__ import annotations -import importlib -import importlib.util +import argparse +import sys import uuid from datetime import datetime, timezone from logging import getLogger from pathlib import Path -from typing import Callable -from microbots.auto_memory.callbacks import ShellCallbackRunner +from microbots.auto_memory.callbacks import CallbackRunner, ShellCallbackRunner from microbots.auto_memory.config import TaskConfig from microbots.auto_memory.errors import ConfigError from microbots.auto_memory.orchestrator import RunSummary, TrainingLoopOrchestrator from microbots.auto_memory.runners.base import AgentRunner +from microbots.auto_memory.runners.writing_bot_runner import WritingBotRunner from microbots.auto_memory.workspace import WorkspaceManager logger = getLogger(__name__) -def _load_runner_class(runner_spec: str, base_dir: Path) -> Callable[..., AgentRunner]: - """Resolve a runner class from a task config ``runner`` string. - - Two forms are supported: - - * **Dotted import path** — ``"pkg.module.ClassName"``. Imported via the - normal import system; the module must be importable (installed package - or on ``sys.path``). - * **File path plus class** — ``"path/to/file.py:ClassName"``. Loaded - directly from disk with :mod:`importlib.util`, so the runner can live - outside the microbots package. Relative file paths are resolved against - *base_dir* (the task YAML's directory). - - Parameters - ---------- - runner_spec : str - The ``runner`` value from the task configuration. - base_dir : Path - Directory used to resolve relative file paths (the task YAML's dir). - - Returns - ------- - Callable[..., AgentRunner] - The resolved runner factory — a class or any callable that accepts - ``model=...`` plus ``runner_params`` and returns an - :class:`~microbots.auto_memory.runners.base.AgentRunner`. - - Raises - ------ - ConfigError - If the spec is malformed, the module/file cannot be imported, the - class is not found, or the resolved attribute is not callable. - """ - if ":" in runner_spec: - # File-path form: ".py:" - file_part, _, cls_name = runner_spec.rpartition(":") - if not file_part or not cls_name: - raise ConfigError( - f"Invalid runner spec '{runner_spec}'; expected 'path/to/file.py:ClassName'" - ) - file_path = Path(file_part) - if not file_path.is_absolute(): - file_path = (base_dir / file_path).resolve() - if not file_path.is_file(): - raise ConfigError(f"Runner file not found: {file_path}") - - module_spec = importlib.util.spec_from_file_location( - "_microbots_user_runner", file_path - ) - if module_spec is None or module_spec.loader is None: - raise ConfigError(f"Cannot load runner module from {file_path}") - module = importlib.util.module_from_spec(module_spec) - try: - module_spec.loader.exec_module(module) - except Exception as exc: # noqa: BLE001 - surface any import-time error - raise ConfigError( - f"Failed to import runner file {file_path}: {exc}" - ) from exc - else: - # Dotted import path form: "pkg.module.ClassName" - module_path, _, cls_name = runner_spec.rpartition(".") - if not module_path or not cls_name: - raise ConfigError( - f"Invalid runner spec '{runner_spec}'; expected " - f"'pkg.module.ClassName' or 'path/to/file.py:ClassName'" - ) - try: - module = importlib.import_module(module_path) - except ImportError as exc: - raise ConfigError( - f"Cannot import runner module '{module_path}': {exc}" - ) from exc - except Exception as exc: # noqa: BLE001 - surface import-time errors - raise ConfigError( - f"Failed to import runner module '{module_path}': {exc}" - ) from exc - - try: - runner_obj = getattr(module, cls_name) - except AttributeError as exc: - raise ConfigError( - f"Runner class '{cls_name}' not found in '{runner_spec}'" - ) from exc - - if not callable(runner_obj): - raise ConfigError( - f"Runner '{cls_name}' in '{runner_spec}' is not callable " - f"(got {type(runner_obj).__name__}); expected a class or factory " - f"that accepts model=... and returns an AgentRunner" - ) - return runner_obj - - def run_from_yaml( yaml_path: str | Path, - workdir: str | Path, + workdir: str | Path | None = None, run_id: str | None = None, *, - model: str, + model: str | None = None, + external_memory_dir: str | Path | None = None, + agent_runner: AgentRunner | None = None, + callback_runner: CallbackRunner | None = None, ) -> RunSummary: """Load a task YAML, wire all components, and run the iteration loop. @@ -139,15 +49,27 @@ def run_from_yaml( ---------- yaml_path : str | Path Path to the task configuration YAML file. - workdir : str | Path - Parent directory that holds the ``runs/`` tree. + workdir : str | Path | None, optional + Parent directory that holds the ``runs/`` tree. Defaults to the YAML + ``workdir`` value, resolved relative to the YAML file's directory. run_id : str | None, optional Identifier for this run. When ``None`` a UTC timestamp plus a short random suffix of the form ``run-YYYYMMDD-HHMMSS-ffffff-`` is generated to avoid collisions. - model : str - Model identifier forwarded to the configured runner (required, - keyword-only — e.g. ``"azure-openai/gpt-4o"``). + model : str | None, optional + Model identifier forwarded to the configured runner. Overrides the + YAML ``model`` value when provided. + external_memory_dir : str | Path | None, optional + If provided, mount this pre-populated directory as the run's memory + directory instead of creating ``/memory/``. Useful for + reusing notes produced by the training loop. Non-feedback files + inside it are never modified. + agent_runner : AgentRunner | None, optional + User-constructed runner for custom agent behavior. Defaults to a + :class:`WritingBotRunner` configured with the resolved model. + callback_runner : CallbackRunner | None, optional + User-provided callback runner. Defaults to :class:`ShellCallbackRunner`, + which executes the callback commands declared in the task YAML. Returns ------- @@ -157,34 +79,44 @@ def run_from_yaml( yaml_path = Path(yaml_path) config = TaskConfig.load_from_yaml(str(yaml_path)) + resolved_model = model or config.model + if agent_runner is None and resolved_model is None: + raise ConfigError( + "A model is required; set 'model' in the task YAML or pass model=..." + ) + + if workdir is None: + configured_workdir = Path(config.workdir) + workdir = ( + configured_workdir + if configured_workdir.is_absolute() + else yaml_path.resolve().parent / configured_workdir + ) + if run_id is None: run_id = _generate_run_id() run_dir = Path(workdir) / "runs" / run_id logger.info("auto_memory: starting run %s at %s", run_id, run_dir) - runner_cls = _load_runner_class(config.runner, base_dir=yaml_path.resolve().parent) - try: - agent_runner: AgentRunner = runner_cls(model=model, **config.runner_params) - except Exception as exc: # noqa: BLE001 - surface construction errors as config errors - raise ConfigError( - f"Failed to construct runner '{config.runner}' with " - f"runner_params={config.runner_params!r}: {exc}" - ) from exc - - # Structural check: the constructed object must satisfy the AgentRunner - # protocol (i.e. expose a run() method). This only verifies method - # presence, not its signature, but catches gross misconfigurations early - # with a clear error instead of failing deep inside the orchestrator. + if agent_runner is None: + assert resolved_model is not None + agent_runner = WritingBotRunner(model=resolved_model) + + # Custom runners must explicitly implement the framework extension point. if not isinstance(agent_runner, AgentRunner): raise ConfigError( - f"Runner '{config.runner}' does not satisfy the AgentRunner " - f"protocol; it must define run(ctx, timeout_s)" + "The provided agent_runner must inherit AgentRunner and implement " + "run(ctx, timeout_s)" ) - logger.info("auto_memory: using runner %s", config.runner) + logger.info("auto_memory: using runner %s", type(agent_runner).__name__) - workspace = WorkspaceManager(run_dir=run_dir) - callback_runner = ShellCallbackRunner() + workspace = WorkspaceManager( + run_dir=run_dir, + external_memory_dir=Path(external_memory_dir) if external_memory_dir else None, + ) + if callback_runner is None: + callback_runner = ShellCallbackRunner() orchestrator = TrainingLoopOrchestrator( config=config, @@ -208,3 +140,54 @@ def _generate_run_id() -> str: """ timestamp = datetime.now(tz=timezone.utc).strftime("%Y%m%d-%H%M%S-%f") return f"run-{timestamp}-{uuid.uuid4().hex[:8]}" + + +def main(argv: list[str] | None = None) -> int: + """Run an auto-memory task from a YAML file. + + Parameters + ---------- + argv : list[str] | None, optional + Arguments to parse. Uses :data:`sys.argv` when omitted. + + Returns + ------- + int + Process exit code, with zero indicating a completed run. + """ + parser = argparse.ArgumentParser( + prog="python -m microbots.auto_memory", + description="Run an iterative auto-memory feedback loop.", + ) + parser.add_argument("yaml_path", type=Path, help="Task YAML file.") + parser.add_argument("--model", help="Override the model declared in YAML.") + parser.add_argument( + "--workdir", type=Path, help="Override the work directory declared in YAML." + ) + parser.add_argument("--run-id", help="Use a fixed run identifier.") + parser.add_argument( + "--external-memory-dir", + type=Path, + help="Reuse an existing memory directory.", + ) + args = parser.parse_args(argv) + + try: + summary = run_from_yaml( + args.yaml_path, + workdir=args.workdir, + run_id=args.run_id, + model=args.model, + external_memory_dir=args.external_memory_dir, + ) + except ConfigError as exc: + print(f"config error: {exc}", file=sys.stderr) + return 2 + + print( + f"auto-memory {summary.final_status.value}: " + f"iterations={summary.iterations_run} elapsed={summary.elapsed_s:.1f}s" + ) + if summary.error_message: + print(f"last error: {summary.error_message}", file=sys.stderr) + return 0 diff --git a/src/microbots/auto_memory/config.py b/src/microbots/auto_memory/config.py index 822d1c25..2b6f7a90 100644 --- a/src/microbots/auto_memory/config.py +++ b/src/microbots/auto_memory/config.py @@ -13,6 +13,29 @@ logger = getLogger(__name__) +DEFAULT_PROMPT_TEMPLATE = """\ +{{ task }} +{% if reference_inputs %} +Reference inputs: +{% for item in reference_inputs %}- {{ item.name }}: {{ item.value }} +{% endfor %}{% endif %} +{% if feedback %} +The previous attempt did not pass validation. Address this feedback: +{{ feedback.summary }} +{% if feedback.root_causes %} +Root causes: +{% for cause in feedback.root_causes %}- {{ cause }} +{% endfor %}{% endif %} +{% if feedback.validator_failures %} +Validator failures: +{% for failure in feedback.validator_failures %}- {{ failure }} +{% endfor %}{% endif %} +{% if feedback.suggested_actions %} +Suggested actions: +{% for action in feedback.suggested_actions %}- {{ action }} +{% endfor %}{% endif %} +{% endif %}""" + @dataclass class TaskConfig: @@ -20,10 +43,12 @@ class TaskConfig: # --- required --- task_definition: str - prompt_template: str callbacks: list[CallbackSpec] # --- optional with defaults --- + prompt_template: str = DEFAULT_PROMPT_TEMPLATE + model: str | None = None + workdir: str = ".auto-memory" reference_inputs: list[ReferenceInput] = field(default_factory=list) output_format: str = "dir" # "file" | "dir" | "stdout" output_path: str = "candidate" # relative to iteration dir @@ -32,19 +57,10 @@ class TaskConfig: per_iteration_timeout: int = 600 # seconds # --- analyzer (LogAnalysisBot) settings --- - analyzer_model: str = "azure-openai/gpt-4o" + analyzer_model: str = "azure-openai/gpt-5" analyzer_max_iterations: int = 20 analyzer_timeout_s: int = 300 - # --- runner selection --- - # Either a dotted import path ("pkg.module.ClassName") or a file-path form - # ("path/to/file.py:ClassName") resolved relative to the task YAML's dir. - # Defaults to the built-in WritingBotRunner so existing configs are - # unaffected. Extra keyword arguments for the runner's constructor go in - # runner_params (model is always injected separately by the CLI). - runner: str = "microbots.auto_memory.runners.writing_bot_runner.WritingBotRunner" - runner_params: dict = field(default_factory=dict) - # ----------------------------------------------------------------------- @classmethod @@ -80,8 +96,17 @@ def load_from_yaml(cls, path: str) -> "TaskConfig": if not isinstance(data, dict): raise ConfigError(f"Expected a YAML mapping at the top level in {path}") + legacy_runner_fields = {"runner", "runner_params"}.intersection(data) + if legacy_runner_fields: + fields = ", ".join(sorted(legacy_runner_fields)) + raise ConfigError( + f"YAML runner configuration ({fields}) is not supported; " + "construct a custom AgentRunner in Python and pass it as " + "run_from_yaml(..., agent_runner=runner)" + ) + # required fields - for required in ("task_definition", "prompt_template"): + for required in ("task_definition",): if required not in data: raise ConfigError(f"Missing required field '{required}' in {path}") @@ -110,27 +135,19 @@ def load_from_yaml(cls, path: str) -> "TaskConfig": config = cls( task_definition=data["task_definition"].strip(), - prompt_template=data["prompt_template"], callbacks=callbacks, + prompt_template=data.get("prompt_template", DEFAULT_PROMPT_TEMPLATE), + model=(str(data["model"]) if data.get("model") is not None else None), + workdir=str(data.get("workdir", ".auto-memory")), reference_inputs=reference_inputs, output_format=data.get("output_format", "dir"), output_path=data.get("output_path", "candidate"), max_iterations=int(data.get("max_iterations", 5)), timeout_min=int(data.get("timeout_min", 60)), per_iteration_timeout=int(data.get("per_iteration_timeout", 600)), - analyzer_model=str(data.get("analyzer_model", "azure-openai/gpt-4o")), + analyzer_model=str(data.get("analyzer_model", "azure-openai/gpt-5")), analyzer_max_iterations=int(data.get("analyzer_max_iterations", 20)), analyzer_timeout_s=int(data.get("analyzer_timeout_s", 300)), - runner=str( - data.get( - "runner", - "microbots.auto_memory.runners.writing_bot_runner.WritingBotRunner", - ) - ), - # Pass through as-is (no dict() coercion) so a non-mapping value - # reaches validate() and produces a clear ConfigError instead of a - # cryptic ValueError/TypeError from dict(). - runner_params=data.get("runner_params", {}) or {}, ) config.validate() return config @@ -147,6 +164,12 @@ def validate(self) -> None: if not self.prompt_template: raise ConfigError("'prompt_template' must not be empty") + if self.model is not None: + self._validate_model(self.model, "model") + + if not self.workdir.strip(): + raise ConfigError("'workdir' must not be empty") + if self.max_iterations < 1: raise ConfigError(f"'max_iterations' must be >= 1, got {self.max_iterations}") @@ -158,24 +181,7 @@ def validate(self) -> None: f"'per_iteration_timeout' must be >= 1, got {self.per_iteration_timeout}" ) - if not self.analyzer_model: - raise ConfigError("'analyzer_model' must not be empty") - - # Mirror MicroBot._validate_model_and_provider so we fail fast at - # config load time instead of deferring to a runtime ValueError - # when LogAnalysisBot is instantiated. - if self.analyzer_model.count("/") != 1: - raise ConfigError( - f"'analyzer_model' must be in the format '/', " - f"got '{self.analyzer_model}'" - ) - provider = self.analyzer_model.split("/", 1)[0] - supported = [e.value for e in ModelProvider] - if provider not in supported: - raise ConfigError( - f"'analyzer_model' has unsupported provider '{provider}'; " - f"expected one of {supported}" - ) + self._validate_model(self.analyzer_model, "analyzer_model") if self.analyzer_max_iterations < 1: raise ConfigError( @@ -221,16 +227,28 @@ def validate(self) -> None: f"reference_input '{ri.name}' must have a non-empty 'value'" ) - if not self.runner or not self.runner.strip(): - raise ConfigError("'runner' must not be empty") + @staticmethod + def _validate_model(model: str, field_name: str) -> None: + """Validate a model identifier using the framework provider contract. - if not isinstance(self.runner_params, dict): + Parameters + ---------- + model : str + Model identifier in ``/`` form. + field_name : str + Configuration field name used in validation errors. + """ + if not model: + raise ConfigError(f"'{field_name}' must not be empty") + if model.count("/") != 1: raise ConfigError( - f"'runner_params' must be a mapping, got {type(self.runner_params).__name__}" + f"'{field_name}' must be in the format '/', " + f"got '{model}'" ) - - if "model" in self.runner_params: + provider = model.split("/", 1)[0] + supported = [e.value for e in ModelProvider] + if provider not in supported: raise ConfigError( - "'runner_params' must not contain 'model'; it is injected " - "automatically by the CLI" + f"'{field_name}' has unsupported provider '{provider}'; " + f"expected one of {supported}" ) diff --git a/src/microbots/auto_memory/loop.py b/src/microbots/auto_memory/loop.py new file mode 100644 index 00000000..e1abb4d1 --- /dev/null +++ b/src/microbots/auto_memory/loop.py @@ -0,0 +1,193 @@ +"""Shared bounded, timed iteration-loop skeleton. + +:class:`~microbots.auto_memory.orchestrator.TrainingLoopOrchestrator` (the +eval loop) and :class:`~microbots.auto_memory.training.orchestrator.TrainingOrchestrator` +(the training loop) both repeatedly invoke a "do one iteration" step, +enforce a total wall-clock timeout, translate an iteration-supplied stop +signal (pass / fail / timeout) into a run-level result, and normalise +unexpected exceptions into an error record. That control flow is +identical between the two loops even though what happens *inside* one +iteration (callbacks + feedback vs. prompt + memory notes) is completely +different. + +This module owns exactly the shared control flow via :func:`run_bounded_loop`. +It knows nothing about bots, callbacks, or memory — those concerns stay in +the concrete orchestrators, which remain separate and independently +configurable. +""" + +from __future__ import annotations + +import time +from dataclasses import dataclass, field +from enum import StrEnum +from logging import getLogger +from typing import Callable, Generic, TypeVar + +logger = getLogger(__name__) + +RecordT = TypeVar("RecordT") + + +class StopReason(StrEnum): + """Why a bounded loop stopped iterating. + + ``SUCCESS`` and ``ITERATION_TIMEOUT`` are only ever produced by the + caller's ``step`` callable; ``LIMIT_REACHED``, ``TOTAL_TIMEOUT``, and + ``ERROR`` (from an uncaught exception) are produced by the loop itself. + """ + + SUCCESS = "success" + LIMIT_REACHED = "limit_reached" + TOTAL_TIMEOUT = "total_timeout" + ITERATION_TIMEOUT = "iteration_timeout" + ERROR = "error" + + +@dataclass +class StepOutcome(Generic[RecordT]): + """Result of one loop step, returned by the caller-supplied ``step`` callable. + + Attributes + ---------- + record : RecordT + Caller-defined per-iteration record; always appended to the run's + record list regardless of ``stop``/``transient_error``. + stop : StopReason | None, optional + When set, the loop stops immediately after this iteration and the + run finishes with this reason. ``None`` means "keep looping". + Ignored when ``transient_error`` is ``True`` and the retry budget + has not yet been exhausted. + transient_error : bool, optional + When ``True``, this iteration counts against + ``max_transient_retries`` instead of stopping outright. Once the + budget is exhausted the loop stops with :attr:`StopReason.ERROR`. + A non-transient iteration resets the consecutive-retry counter. + """ + + record: RecordT + stop: StopReason | None = None + transient_error: bool = False + + +@dataclass +class LoopResult(Generic[RecordT]): + """Outcome of a full :func:`run_bounded_loop` run.""" + + stop_reason: StopReason + iterations_run: int + records: list[RecordT] = field(default_factory=list) + elapsed_s: float = 0.0 + error_message: str | None = None + + +def run_bounded_loop( + *, + max_iterations: int, + total_timeout_s: float, + step: Callable[[int], StepOutcome[RecordT]], + catch_exceptions: tuple[type[BaseException], ...] = (Exception,), + on_exception: Callable[[int, BaseException], RecordT] | None = None, + max_transient_retries: int = 0, +) -> LoopResult[RecordT]: + """Drive ``step`` for up to ``max_iterations``, honouring a total timeout. + + Before every iteration the elapsed wall-clock time is compared against + ``total_timeout_s``; if it has been exceeded the loop stops with + :attr:`StopReason.TOTAL_TIMEOUT` *without* running that iteration + (``iterations_run`` reflects only completed iterations). + + If ``step`` raises one of ``catch_exceptions``, the loop stops with + :attr:`StopReason.ERROR`; ``on_exception`` (if given) is called to build + a final record for that iteration, and ``error_message`` is set to + ``str(exc)``. Exceptions not in ``catch_exceptions`` propagate to the + caller uncaught, exactly as if this loop were not there. + + A returned :class:`StepOutcome` with ``transient_error=True`` counts + against ``max_transient_retries``: while the consecutive count is at or + below the budget the loop continues to the next iteration; once + exceeded it stops with :attr:`StopReason.ERROR` (``error_message`` is + left ``None`` — callers typically derive a message from the last + record). Any non-transient iteration resets the counter. + + Otherwise, ``step``'s ``stop`` value (if not ``None``) ends the loop + immediately with that reason; ``None`` continues to the next iteration. + If ``max_iterations`` is exhausted without a stop, the loop ends with + :attr:`StopReason.LIMIT_REACHED`. + + Parameters + ---------- + max_iterations : int + Maximum number of iterations to run. + total_timeout_s : float + Total wall-clock budget for the whole loop, in seconds. A value + ``<= 0`` disables the total-timeout check. + step : Callable[[int], StepOutcome[RecordT]] + Runs one iteration (given its zero-based index) and reports its + outcome. May raise; see ``catch_exceptions``. + catch_exceptions : tuple[type[BaseException], ...], optional + Exception types that :func:`run_bounded_loop` intercepts and turns + into a :attr:`StopReason.ERROR` result. Defaults to ``(Exception,)``. + Pass a narrower tuple (e.g. ``(AgentError,)``) to let unrelated bugs + propagate uncaught. + on_exception : Callable[[int, BaseException], RecordT] | None, optional + Builds the final record to append when ``step`` raises a caught + exception. If omitted, no record is appended for that iteration. + max_transient_retries : int, optional + Number of consecutive ``transient_error=True`` outcomes to tolerate + before stopping. Defaults to ``0`` (no retries — the first + transient error stops the loop). + + Returns + ------- + LoopResult[RecordT] + The accumulated records, stop reason, elapsed time, and (when + applicable) error message. + """ + start = time.monotonic() + records: list[RecordT] = [] + consecutive_transient = 0 + + for idx in range(max_iterations): + elapsed = time.monotonic() - start + if total_timeout_s > 0 and elapsed >= total_timeout_s: + logger.info( + "Bounded loop: total timeout reached after %.1fs (limit %.0fs)", + elapsed, + total_timeout_s, + ) + return LoopResult(StopReason.TOTAL_TIMEOUT, idx, records, elapsed) + + try: + outcome = step(idx) + except catch_exceptions as exc: # noqa: BLE001 - intentionally caller-scoped + elapsed = time.monotonic() - start + if on_exception is not None: + records.append(on_exception(idx, exc)) + logger.error("Bounded loop: iteration %d raised %s", idx, exc) + return LoopResult( + StopReason.ERROR, idx + 1, records, elapsed, error_message=str(exc) + ) + + records.append(outcome.record) + elapsed = time.monotonic() - start + + if outcome.transient_error: + consecutive_transient += 1 + if consecutive_transient <= max_transient_retries: + logger.warning( + "Bounded loop: transient error on iteration %d (retry %d/%d), continuing", + idx, + consecutive_transient, + max_transient_retries, + ) + continue + return LoopResult(StopReason.ERROR, idx + 1, records, elapsed) + + consecutive_transient = 0 + + if outcome.stop is not None: + return LoopResult(outcome.stop, idx + 1, records, elapsed) + + elapsed = time.monotonic() - start + return LoopResult(StopReason.LIMIT_REACHED, max_iterations, records, elapsed) diff --git a/src/microbots/auto_memory/memory.py b/src/microbots/auto_memory/memory.py index baba7227..e4c36937 100644 --- a/src/microbots/auto_memory/memory.py +++ b/src/microbots/auto_memory/memory.py @@ -40,8 +40,20 @@ def __init__(self) -> None: # ------------------------------------------------------------------ # Lifecycle - def mount(self, run_dir: Path, *, resume: bool = False) -> None: - """Attach the store to *run_dir/memory/*. + def mount( + self, + run_dir: Path, + *, + resume: bool = False, + external_memory_dir: Path | None = None, + ) -> None: + """Attach the store to a memory directory. + + By default the store lives at ``/memory/`` — a fresh dir + owned by this run. Pass ``external_memory_dir`` to reuse an + already-populated memory dir (e.g. notes produced by the training + loop). Only the feedback JSONL inside it is managed by the store; + any other files (e.g. ``*.md`` notes) are left untouched. ``resume=False`` (default): the feedback file is wiped so the new run starts with a clean slate. ``resume=True``: if the file already @@ -54,13 +66,20 @@ def mount(self, run_dir: Path, *, resume: bool = False) -> None: Root directory of the auto_memory run. resume : bool, optional When ``True``, preserve any existing feedback file. + external_memory_dir : Path | None, optional + If provided, use this directory as the memory root instead of + ``/memory/``. The directory is created if missing. + Non-feedback files (e.g. trained notes) are never modified. Raises ------ MemoryStoreError If the memory directory cannot be created. """ - memory_dir = run_dir / "memory" + memory_dir = ( + Path(external_memory_dir) if external_memory_dir is not None + else run_dir / "memory" + ) try: memory_dir.mkdir(parents=True, exist_ok=True) except OSError as exc: @@ -75,7 +94,10 @@ def mount(self, run_dir: Path, *, resume: bool = False) -> None: self.clear() logger.debug( - "MemoryStore mounted at %s (resume=%s)", memory_dir, resume + "MemoryStore mounted at %s (resume=%s, external=%s)", + memory_dir, + resume, + external_memory_dir is not None, ) # ------------------------------------------------------------------ diff --git a/src/microbots/auto_memory/orchestrator.py b/src/microbots/auto_memory/orchestrator.py index b70a9771..079e6393 100644 --- a/src/microbots/auto_memory/orchestrator.py +++ b/src/microbots/auto_memory/orchestrator.py @@ -2,7 +2,6 @@ from __future__ import annotations -import time from dataclasses import dataclass, field from logging import getLogger from pathlib import Path @@ -13,11 +12,21 @@ from microbots.auto_memory.context import build_iteration_context from microbots.auto_memory.data_models import Feedback, FinalStatus, IterationStatus from microbots.auto_memory.errors import AgentError +from microbots.auto_memory.loop import StepOutcome, StopReason, run_bounded_loop from microbots.auto_memory.runners.base import AgentResult, AgentRunner, IterationContext from microbots.auto_memory.workspace import WorkspaceManager logger = getLogger(__name__) +# Maps the generic loop's stop reason onto this loop's own FinalStatus enum. +_STOP_REASON_TO_FINAL_STATUS: dict[StopReason, FinalStatus] = { + StopReason.SUCCESS: FinalStatus.PASSED, + StopReason.LIMIT_REACHED: FinalStatus.LIMIT_REACHED, + StopReason.TOTAL_TIMEOUT: FinalStatus.TIMEOUT, + StopReason.ITERATION_TIMEOUT: FinalStatus.TIMEOUT, + StopReason.ERROR: FinalStatus.ERROR, +} + # --------------------------------------------------------------------------- # RunSummary @@ -82,7 +91,7 @@ class TrainingLoopOrchestrator: config : TaskConfig Task configuration (max_iterations, timeout_min, etc.). agent_runner : AgentRunner - Structural-protocol object that executes one agent iteration. + Runner object that executes one agent iteration. callback_runner : CallbackRunner Runs validation callbacks against the agent's output. workspace : WorkspaceManager @@ -108,7 +117,7 @@ def __init__( config : TaskConfig Task configuration (max_iterations, timeout_min, etc.). agent_runner : AgentRunner - Structural-protocol object that executes one agent iteration. + Runner object that executes one agent iteration. callback_runner : CallbackRunner Runs validation callbacks against the agent's output. workspace : WorkspaceManager @@ -145,118 +154,101 @@ def run(self) -> RunSummary: """ self._workspace.prepare() - start_time = time.monotonic() - timeout_s = self._config.timeout_min * 60 - records: list[IterationRecord] = [] last_feedback: Feedback | None = None - consecutive_errors = 0 - - for iteration_idx in range(self._config.max_iterations): - elapsed = time.monotonic() - start_time - # --- total timeout check --- - if elapsed >= timeout_s: - logger.info( - "Orchestrator: total timeout reached after %.1fs (limit %ds)", - elapsed, - timeout_s, - ) - return RunSummary( - final_status=FinalStatus.TIMEOUT, - iterations_run=iteration_idx, - iteration_records=records, - elapsed_s=elapsed, - ) - - # --- run one iteration --- - try: - record = self.run_iteration( - iteration_idx=iteration_idx, - feedback=last_feedback, - ) - except AgentError as exc: - elapsed = time.monotonic() - start_time - logger.error( - "Orchestrator: AgentError on iteration %d: %s", iteration_idx, exc - ) - return RunSummary( - final_status=FinalStatus.ERROR, - iterations_run=iteration_idx + 1, - iteration_records=records, - elapsed_s=elapsed, - error_message=str(exc), - ) - - records.append(record) + def step(iteration_idx: int) -> StepOutcome[IterationRecord]: + """Run one iteration and translate its status into a loop outcome. + + Parameters + ---------- + iteration_idx : int + Zero-based index of the iteration to run. + + Returns + ------- + StepOutcome[IterationRecord] + The iteration's record plus the loop-control signal derived + from its status (retry, stop, or continue). + """ + nonlocal last_feedback + record = self.run_iteration( + iteration_idx=iteration_idx, feedback=last_feedback + ) if record.status == IterationStatus.ERROR: - consecutive_errors += 1 - if consecutive_errors <= self._max_agent_retries: - logger.warning( - "Orchestrator: transient agent error on iteration %d " - "(retry %d/%d), continuing", - iteration_idx, - consecutive_errors, - self._max_agent_retries, - ) - continue - elapsed = time.monotonic() - start_time - error_msg = ( - record.error - or f"Agent returned ERROR on iteration {iteration_idx}" - ) - logger.error( - "Orchestrator: iteration %d returned ERROR (%d consecutive)", - iteration_idx, - consecutive_errors, - ) - return RunSummary( - final_status=FinalStatus.ERROR, - iterations_run=iteration_idx + 1, - iteration_records=records, - elapsed_s=elapsed, - error_message=error_msg, - ) - - consecutive_errors = 0 + return StepOutcome(record=record, transient_error=True) if record.status == IterationStatus.TIMEOUT: - elapsed = time.monotonic() - start_time logger.info( "Orchestrator: per-iteration timeout on iteration %d", iteration_idx ) - return RunSummary( - final_status=FinalStatus.TIMEOUT, - iterations_run=iteration_idx + 1, - iteration_records=records, - elapsed_s=elapsed, - ) + return StepOutcome(record=record, stop=StopReason.ITERATION_TIMEOUT) if record.status == IterationStatus.PASSED: - elapsed = time.monotonic() - start_time - logger.info( - "Orchestrator: PASSED on iteration %d (%.1fs)", iteration_idx, elapsed - ) - return RunSummary( - final_status=FinalStatus.PASSED, - iterations_run=iteration_idx + 1, - iteration_records=records, - elapsed_s=elapsed, - ) + logger.info("Orchestrator: PASSED on iteration %d", iteration_idx) + return StepOutcome(record=record, stop=StopReason.SUCCESS) # FAILED — persist feedback and continue last_feedback = record.feedback + return StepOutcome(record=record) + + def on_exception(iteration_idx: int, exc: BaseException) -> IterationRecord: + """Build the final iteration record when the agent raises ``AgentError``. + + Parameters + ---------- + iteration_idx : int + Zero-based index of the iteration that raised. + exc : BaseException + The caught ``AgentError`` instance. + + Returns + ------- + IterationRecord + An ``ERROR`` record carrying the exception message. + """ + logger.error( + "Orchestrator: AgentError on iteration %d: %s", iteration_idx, exc + ) + return IterationRecord( + idx=iteration_idx, status=IterationStatus.ERROR, error=str(exc) + ) + + result = run_bounded_loop( + max_iterations=self._config.max_iterations, + total_timeout_s=self._config.timeout_min * 60, + step=step, + catch_exceptions=(AgentError,), + on_exception=on_exception, + max_transient_retries=self._max_agent_retries, + ) + + final_status = _STOP_REASON_TO_FINAL_STATUS[result.stop_reason] + + error_message = result.error_message + if ( + final_status == FinalStatus.ERROR + and error_message is None + and result.records + ): + last_record = result.records[-1] + error_message = ( + last_record.error + or f"Agent returned ERROR on iteration {last_record.idx}" + ) - # All iterations exhausted without a pass - elapsed = time.monotonic() - start_time logger.info( - "Orchestrator: limit reached after %d iteration(s)", self._config.max_iterations + "Orchestrator: finished status=%s iterations=%d elapsed=%.1fs", + final_status, + result.iterations_run, + result.elapsed_s, ) return RunSummary( - final_status=FinalStatus.LIMIT_REACHED, - iterations_run=self._config.max_iterations, - iteration_records=records, - elapsed_s=elapsed, + final_status=final_status, + iterations_run=result.iterations_run, + iteration_records=result.records, + elapsed_s=result.elapsed_s, + error_message=error_message, ) def run_iteration( diff --git a/src/microbots/auto_memory/runners/__init__.py b/src/microbots/auto_memory/runners/__init__.py index 5aa610cc..cb6d5552 100644 --- a/src/microbots/auto_memory/runners/__init__.py +++ b/src/microbots/auto_memory/runners/__init__.py @@ -8,8 +8,7 @@ * ``IterationContext`` — immutable context passed to every :class:`AgentRunner`. * ``AgentResult`` — normalised result returned by every runner. -* ``AgentRunner`` — structural protocol satisfied by any runner with a matching - ``run`` method. +* ``AgentRunner`` — abstract base class implemented by every runner. """ from microbots.auto_memory.runners.base import AgentResult, AgentRunner, IterationContext diff --git a/src/microbots/auto_memory/runners/base.py b/src/microbots/auto_memory/runners/base.py index 93e4db57..027c07af 100644 --- a/src/microbots/auto_memory/runners/base.py +++ b/src/microbots/auto_memory/runners/base.py @@ -2,8 +2,8 @@ from __future__ import annotations +from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import Protocol, runtime_checkable from microbots.auto_memory.data_models import IterationStatus @@ -59,10 +59,10 @@ class AgentResult: log_path: str | None = None -@runtime_checkable -class AgentRunner(Protocol): - """Structural protocol satisfied by any object with a matching ``run`` method.""" +class AgentRunner(ABC): + """Abstract base class for one auto-memory agent invocation.""" + @abstractmethod def run(self, ctx: IterationContext, timeout_s: int) -> AgentResult: """Run the agent described by *ctx* and return a normalised result. diff --git a/src/microbots/auto_memory/runners/writing_bot_runner.py b/src/microbots/auto_memory/runners/writing_bot_runner.py index b413616b..1181f6c3 100644 --- a/src/microbots/auto_memory/runners/writing_bot_runner.py +++ b/src/microbots/auto_memory/runners/writing_bot_runner.py @@ -15,10 +15,10 @@ _TIMEOUT_ERROR_PREFIX = "Timeout of " -class WritingBotRunner: +class WritingBotRunner(AgentRunner): """Runs a :class:`~microbots.bot.WritingBot.WritingBot` for one iteration. - Satisfies the :class:`AgentRunner` protocol structurally. + Implements the :class:`AgentRunner` abstract base class. Parameters ---------- diff --git a/src/microbots/auto_memory/training/__init__.py b/src/microbots/auto_memory/training/__init__.py index ad0a6abc..919e0033 100644 --- a/src/microbots/auto_memory/training/__init__.py +++ b/src/microbots/auto_memory/training/__init__.py @@ -21,6 +21,7 @@ from microbots.auto_memory.training.config import TrainingConfig from microbots.auto_memory.training.training_source import TrainingSource from microbots.auto_memory.training.orchestrator import ( + TrainingFinalStatus, TrainingIterationRecord, TrainingOrchestrator, TrainingSummary, @@ -36,6 +37,7 @@ "TrainingOrchestrator", "TrainingSummary", "TrainingIterationRecord", + "TrainingFinalStatus", "run_training", "run_training_from_yaml", ] diff --git a/src/microbots/auto_memory/training/orchestrator.py b/src/microbots/auto_memory/training/orchestrator.py index ce279873..b77bdb56 100644 --- a/src/microbots/auto_memory/training/orchestrator.py +++ b/src/microbots/auto_memory/training/orchestrator.py @@ -24,10 +24,13 @@ import time from dataclasses import asdict, dataclass, field from datetime import datetime, timezone +from enum import StrEnum from logging import getLogger from pathlib import Path +from microbots.auto_memory.data_models import IterationStatus from microbots.auto_memory.errors import ConfigError +from microbots.auto_memory.loop import StepOutcome, StopReason, run_bounded_loop from microbots.auto_memory.training.config import TrainingConfig from microbots.auto_memory.training.runner import ( LearningRunner, @@ -37,12 +40,29 @@ logger = getLogger(__name__) +class TrainingFinalStatus(StrEnum): + """Overall status of a completed training run.""" + + COMPLETED = "completed" + TIMEOUT = "timeout" + ERROR = "error" + + +# Maps the generic loop's stop reason onto this loop's own TrainingFinalStatus. +_STOP_REASON_TO_FINAL_STATUS: dict[StopReason, TrainingFinalStatus] = { + StopReason.LIMIT_REACHED: TrainingFinalStatus.COMPLETED, + StopReason.TOTAL_TIMEOUT: TrainingFinalStatus.TIMEOUT, + StopReason.ITERATION_TIMEOUT: TrainingFinalStatus.TIMEOUT, + StopReason.ERROR: TrainingFinalStatus.ERROR, +} + + @dataclass class TrainingIterationRecord: """Per-iteration record persisted to ``training_run.jsonl``.""" idx: int - status: str + status: IterationStatus elapsed_s: float error: str | None = None @@ -51,7 +71,7 @@ class TrainingIterationRecord: class TrainingSummary: """Summary of one completed training run.""" - final_status: str # "completed" | "timeout" | "error" + final_status: TrainingFinalStatus iterations_run: int iteration_records: list[TrainingIterationRecord] = field(default_factory=list) elapsed_s: float = 0.0 @@ -121,86 +141,101 @@ def run(self) -> TrainingSummary: agents_md = self._config.read_agents_md() total = self._config.iterations - total_budget_s = self._config.total_timeout_min * 60 - started = time.monotonic() - records: list[TrainingIterationRecord] = [] - - for idx in range(total): - elapsed = time.monotonic() - started - if total_budget_s and elapsed >= total_budget_s: - logger.info( - "TrainingOrchestrator: total timeout reached after %.1fs (limit %ds)", - elapsed, - total_budget_s, - ) - return self._finish("timeout", records, elapsed, error=None) - + last_iter_started = 0.0 + + def step(idx: int) -> StepOutcome[TrainingIterationRecord]: + """Run one training iteration and log its outcome. + + Parameters + ---------- + idx : int + Zero-based index of the iteration to run. + + Returns + ------- + StepOutcome[TrainingIterationRecord] + The iteration's record plus a stop signal when the runner + itself reports a per-iteration timeout. + """ + nonlocal last_iter_started prompt = agents_md + _ITER_HEADER.format( - idx=idx, - total=total, - source=resolved_source, + idx=idx, total=total, source=resolved_source ) - logger.info( "TrainingOrchestrator: iteration %d/%d starting", idx + 1, total ) - iter_started = time.monotonic() - try: - result: TrainingIterationResult = runner.run( - prompt, timeout_s=self._config.per_iteration_timeout - ) - except Exception as exc: # noqa: BLE001 - surface as ERROR record - iter_elapsed = time.monotonic() - iter_started - record = TrainingIterationRecord( - idx=idx, - status="error", - elapsed_s=iter_elapsed, - error=f"{type(exc).__name__}: {exc}", - ) - records.append(record) - self._append_log(record) - logger.exception( - "TrainingOrchestrator: iteration %d raised %s", - idx, - type(exc).__name__, - ) - return self._finish( - "error", - records, - time.monotonic() - started, - error=record.error, - ) - - iter_elapsed = time.monotonic() - iter_started + last_iter_started = time.monotonic() + result: TrainingIterationResult = runner.run( + prompt, timeout_s=self._config.per_iteration_timeout + ) + iter_elapsed = time.monotonic() - last_iter_started record = TrainingIterationRecord( idx=idx, status=result.status, elapsed_s=iter_elapsed, error=result.error, ) - records.append(record) self._append_log(record) - logger.info( "TrainingOrchestrator: iteration %d finished status=%s in %.1fs", idx, result.status, iter_elapsed, ) + stop = ( + StopReason.ITERATION_TIMEOUT + if result.status == IterationStatus.TIMEOUT + else None + ) + return StepOutcome(record=record, stop=stop) + + def on_exception(idx: int, exc: BaseException) -> TrainingIterationRecord: + """Build the final iteration record when the runner raises. + + Parameters + ---------- + idx : int + Zero-based index of the iteration that raised. + exc : BaseException + The caught exception instance. + + Returns + ------- + TrainingIterationRecord + An ``ERROR`` record carrying the exception type and message. + """ + record = TrainingIterationRecord( + idx=idx, + status=IterationStatus.ERROR, + elapsed_s=time.monotonic() - last_iter_started, + error=f"{type(exc).__name__}: {exc}", + ) + self._append_log(record) + logger.exception( + "TrainingOrchestrator: iteration %d raised %s", idx, type(exc).__name__ + ) + return record + + result = run_bounded_loop( + max_iterations=total, + total_timeout_s=self._config.total_timeout_min * 60, + step=step, + catch_exceptions=(Exception,), + on_exception=on_exception, + ) - if result.status == "timeout": - return self._finish( - "timeout", - records, - time.monotonic() - started, - error=result.error, - ) + final_status = _STOP_REASON_TO_FINAL_STATUS[result.stop_reason] + + # Prefer the last record's own error (richer / iteration-scoped) + # over the loop's generic error_message when both are available. + error_message = ( + result.records[-1].error if result.records else result.error_message + ) + if error_message is None: + error_message = result.error_message return self._finish( - "completed", - records, - time.monotonic() - started, - error=None, + final_status, result.records, result.elapsed_s, error=error_message ) def _prepare_workdir(self) -> None: @@ -262,7 +297,7 @@ def _append_log(self, record: TrainingIterationRecord) -> None: def _finish( self, - final_status: str, + final_status: TrainingFinalStatus, records: list[TrainingIterationRecord], elapsed_s: float, error: str | None, @@ -271,7 +306,7 @@ def _finish( Parameters ---------- - final_status : str + final_status : TrainingFinalStatus Overall completion status. records : list[TrainingIterationRecord] Iteration records accumulated during the run. diff --git a/src/microbots/auto_memory/training/runner.py b/src/microbots/auto_memory/training/runner.py index 86962262..9464e72c 100644 --- a/src/microbots/auto_memory/training/runner.py +++ b/src/microbots/auto_memory/training/runner.py @@ -19,6 +19,7 @@ from logging import getLogger from pathlib import Path +from microbots.auto_memory.data_models import IterationStatus from microbots.bot.ReadingBot import ReadingBot from microbots.MicroBot import BotRunResult from microbots.tools.tool_definitions.memory_tool import MemoryTool @@ -34,15 +35,15 @@ class TrainingIterationResult: Attributes ---------- - status : str - One of ``"passed"``, ``"timeout"``, ``"error"``. + status : IterationStatus + One of ``PASSED``, ``TIMEOUT``, ``ERROR``. output : str | None Bot's final answer on success, else ``None``. error : str | None Error description on failure, else ``None``. """ - status: str + status: IterationStatus output: str | None error: str | None @@ -142,14 +143,14 @@ def _map(bot_result: BotRunResult) -> TrainingIterationResult: """ if bot_result.status: return TrainingIterationResult( - status="passed", output=bot_result.result, error=None + status=IterationStatus.PASSED, output=bot_result.result, error=None ) error = bot_result.error or "Unknown error" if error.startswith(_TIMEOUT_PREFIX): return TrainingIterationResult( - status="timeout", output=None, error=error + status=IterationStatus.TIMEOUT, output=None, error=error ) return TrainingIterationResult( - status="error", output=None, error=error + status=IterationStatus.ERROR, output=None, error=error ) diff --git a/src/microbots/auto_memory/workspace.py b/src/microbots/auto_memory/workspace.py index 58bc90b6..b83da498 100644 --- a/src/microbots/auto_memory/workspace.py +++ b/src/microbots/auto_memory/workspace.py @@ -47,10 +47,23 @@ class WorkspaceManager: wm.memory.append_feedback(feedback) wm.cleanup() # strip __pycache__ etc. + + Reusing a pre-populated memory directory (e.g. notes produced by the + training loop) is opt-in via ``external_memory_dir``:: + + wm = WorkspaceManager( + run_dir=Path("runs/my_task"), + external_memory_dir=Path(".memories/pytest"), + ) + + When set, the memory directory lives at ``external_memory_dir`` instead + of ``/memory/``. Non-feedback files in the external directory + are never modified. """ run_dir: Path memory: MemoryStore = field(default_factory=MemoryStore) + external_memory_dir: Path | None = None # Set by prepare(); not part of the constructor signature. _iterations_dir: Path = field(init=False, repr=False) @@ -110,7 +123,11 @@ def prepare(self, *, resume: bool = False) -> None: self._iteration_count = 0 self._write_meta() - self.memory.mount(self.run_dir, resume=resume) + self.memory.mount( + self.run_dir, + resume=resume, + external_memory_dir=self.external_memory_dir, + ) def reset(self) -> None: """Wipe *run_dir* and create a fresh layout. diff --git a/test/auto_memory/runners/test_writing_bot_runner.py b/test/auto_memory/runners/test_writing_bot_runner.py index e327ef08..8be6cd21 100644 --- a/test/auto_memory/runners/test_writing_bot_runner.py +++ b/test/auto_memory/runners/test_writing_bot_runner.py @@ -4,7 +4,7 @@ from unittest.mock import MagicMock, patch from microbots.auto_memory.data_models import IterationStatus -from microbots.auto_memory.runners import AgentResult, IterationContext +from microbots.auto_memory.runners import AgentResult, AgentRunner, IterationContext from microbots.auto_memory.runners.writing_bot_runner import WritingBotRunner from microbots.MicroBot import BotRunResult @@ -23,6 +23,15 @@ def _make_ctx(memory_dir: str, task: str = _TASK) -> IterationContext: ) +@pytest.mark.unit +def test_agent_runner_requires_run_implementation() -> None: + class IncompleteRunner(AgentRunner): + pass + + with pytest.raises(TypeError, match="abstract method.*run"): + IncompleteRunner() + + # --------------------------------------------------------------------------- # WritingBot construction contract # --------------------------------------------------------------------------- diff --git a/test/auto_memory/test_cli.py b/test/auto_memory/test_cli.py index 6c333449..5d12b613 100644 --- a/test/auto_memory/test_cli.py +++ b/test/auto_memory/test_cli.py @@ -3,16 +3,19 @@ from __future__ import annotations import textwrap +import runpy from pathlib import Path from unittest.mock import MagicMock, patch import pytest -from microbots.auto_memory import run_from_yaml -from microbots.auto_memory.cli import _load_runner_class -from microbots.auto_memory.data_models import FinalStatus +from microbots.auto_memory import CallbackRunner, run_from_yaml +from microbots.auto_memory.callbacks import CallbackResult +from microbots.auto_memory.cli import main +from microbots.auto_memory.data_models import FinalStatus, IterationStatus from microbots.auto_memory.errors import ConfigError from microbots.auto_memory.orchestrator import RunSummary +from microbots.auto_memory.runners.base import AgentResult, AgentRunner from microbots.MicroBot import BotRunResult _MODEL = "azure-openai/gpt-4o" @@ -84,6 +87,47 @@ def _mock_log_analysis_bot(result: str = "diagnosis narrative"): @pytest.mark.unit class TestRunFromYamlEndToEnd: + def test_yaml_supplies_model_and_default_workdir(self, tmp_path): + yaml_path = _write_yaml( + tmp_path, + "model: azure-openai/gpt-4o\n" + _TASK_YAML, + ) + bot_patch, _ = _mock_writing_bot() + + with bot_patch, patch( + "microbots.auto_memory.runners.writing_bot_runner.MemoryTool" + ): + summary = run_from_yaml(yaml_path, run_id="yaml-only") + + assert summary.final_status == FinalStatus.PASSED + assert (tmp_path / ".auto-memory" / "runs" / "yaml-only").is_dir() + + def test_explicit_model_and_workdir_override_yaml(self, tmp_path): + yaml_path = _write_yaml( + tmp_path, + "model: azure-openai/from-yaml\nworkdir: yaml-work\n" + _TASK_YAML, + ) + bot_patch, _ = _mock_writing_bot() + explicit_workdir = tmp_path / "explicit-work" + + with bot_patch as writing_bot, patch( + "microbots.auto_memory.runners.writing_bot_runner.MemoryTool" + ): + run_from_yaml( + yaml_path, + explicit_workdir, + run_id="overrides", + model=_MODEL, + ) + + assert writing_bot.call_args.kwargs["model"] == _MODEL + assert (explicit_workdir / "runs" / "overrides").is_dir() + assert not (tmp_path / "yaml-work").exists() + + def test_requires_model_in_yaml_or_argument(self, tmp_path): + with pytest.raises(ConfigError, match="model is required"): + run_from_yaml(_write_yaml(tmp_path)) + def test_returns_run_summary(self, tmp_path): yaml_path = _write_yaml(tmp_path) workdir = tmp_path / "workdir" @@ -115,6 +159,41 @@ def test_final_status_passed_when_callbacks_pass(self, tmp_path): assert summary.error_message is None assert len(summary.iteration_records) == 1 + def test_uses_user_callback_runner(self, tmp_path): + yaml_path = _write_yaml(tmp_path) + workdir = tmp_path / "workdir" + bot_patch, _ = _mock_writing_bot() + + class PassingCallbacks(CallbackRunner): + def run_all(self, specs, logs_dir, candidate_path): + return [ + CallbackResult( + spec=spec, + return_code=0, + stdout_path=logs_dir / f"{spec.name}.stdout", + stderr_path=logs_dir / f"{spec.name}.stderr", + passed=True, + ) + for spec in specs + ] + + callback_runner = PassingCallbacks() + with bot_patch, patch( + "microbots.auto_memory.runners.writing_bot_runner.MemoryTool" + ), patch( + "microbots.auto_memory.cli.ShellCallbackRunner" + ) as shell_callback_runner: + summary = run_from_yaml( + yaml_path, + workdir, + run_id="custom-callbacks", + model=_MODEL, + callback_runner=callback_runner, + ) + + assert summary.final_status == FinalStatus.PASSED + shell_callback_runner.assert_not_called() + def test_disk_layout_created(self, tmp_path): yaml_path = _write_yaml(tmp_path) workdir = tmp_path / "workdir" @@ -193,172 +272,107 @@ def test_failing_callbacks_persist_feedback_and_reach_limit(self, tmp_path): assert (run_dir / "iterations" / "iter_01" / "candidate").is_dir() -# --------------------------------------------------------------------------- -# Runner construction guards (run_from_yaml) -# --------------------------------------------------------------------------- - -_GOOD_RUNNER_SRC = textwrap.dedent("""\ - class R: - def __init__(self, model): - self.model = model - - def run(self, ctx, timeout_s): - return None -""") - -_NO_RUN_RUNNER_SRC = textwrap.dedent("""\ - class R: - def __init__(self, model): - self.model = model -""") - - -def _write_custom_runner_yaml(tmp_path: Path, runner_src: str, runner_params: str) -> Path: - (tmp_path / "custom_runner.py").write_text(runner_src) - yaml = textwrap.dedent(f"""\ - task_definition: do a thing - prompt_template: "Goal: {{{{ task }}}}" - runner: ./custom_runner.py:R - runner_params: - {runner_params} - callbacks: - - name: always_ok - command: 'true' - max_iterations: 1 - timeout_min: 1 - per_iteration_timeout: 30 - """) - p = tmp_path / "task.yml" - p.write_text(yaml) - return p - - @pytest.mark.unit -class TestRunnerConstructionGuards: - def test_bad_runner_params_raises_config_error(self, tmp_path): - """runner_params that don't match __init__ surface as ConfigError.""" - yaml_path = _write_custom_runner_yaml( - tmp_path, _GOOD_RUNNER_SRC, runner_params=" unexpected_kwarg: 1" - ) - with pytest.raises(ConfigError, match="Failed to construct runner"): - run_from_yaml(str(yaml_path), str(tmp_path / "wd"), model=_MODEL) - - def test_runner_missing_run_raises_config_error(self, tmp_path): - """A runner without run() fails the AgentRunner protocol check.""" - yaml_path = _write_custom_runner_yaml( - tmp_path, _NO_RUN_RUNNER_SRC, runner_params=" {}" +class TestAgentRunnerInjection: + def test_uses_user_constructed_runner(self, tmp_path): + class CustomRunner(AgentRunner): + def __init__(self): + self.calls = [] + + def run(self, ctx, timeout_s): + self.calls.append((ctx, timeout_s)) + return AgentResult(IterationStatus.PASSED, "done", None) + + runner = CustomRunner() + summary = run_from_yaml( + _write_yaml(tmp_path), + tmp_path / "workdir", + run_id="custom-runner", + agent_runner=runner, ) - with pytest.raises(ConfigError, match="AgentRunner"): - run_from_yaml(str(yaml_path), str(tmp_path / "wd"), model=_MODEL) - - -# --------------------------------------------------------------------------- -# Runner resolution (_load_runner_class) -# --------------------------------------------------------------------------- -_RUNNER_FILE_SRC = textwrap.dedent("""\ - class MyRunner: - def __init__(self, model, **kwargs): - self.model = model - self.kwargs = kwargs + assert summary.final_status == FinalStatus.PASSED + assert len(runner.calls) == 1 + assert runner.calls[0][0].task == "Goal: Write a hello message to /memories/hello.txt" - def run(self, ctx, timeout_s): - return None -""") + def test_rejects_object_without_runner_protocol(self, tmp_path): + with pytest.raises(ConfigError, match="AgentRunner"): + run_from_yaml( + _write_yaml(tmp_path), + tmp_path / "workdir", + agent_runner=object(), # type: ignore[arg-type] + ) @pytest.mark.unit -class TestLoadRunnerClass: - def _write_runner(self, tmp_path: Path, name: str = "myrunner.py") -> Path: - p = tmp_path / name - p.write_text(_RUNNER_FILE_SRC) - return p - - # --- file-path form ------------------------------------------------- - def test_file_path_form_loads_class(self, tmp_path): - self._write_runner(tmp_path) - cls = _load_runner_class("myrunner.py:MyRunner", base_dir=tmp_path) - assert cls.__name__ == "MyRunner" - instance = cls(model=_MODEL, repo_url="x") - assert instance.model == _MODEL - assert instance.kwargs == {"repo_url": "x"} - - def test_file_path_form_absolute(self, tmp_path): - runner = self._write_runner(tmp_path) - cls = _load_runner_class(f"{runner}:MyRunner", base_dir=Path("/nonexistent")) - assert cls.__name__ == "MyRunner" - - def test_file_path_missing_class_name(self, tmp_path): - self._write_runner(tmp_path) - with pytest.raises(ConfigError, match="expected 'path/to/file.py:ClassName'"): - _load_runner_class("myrunner.py:", base_dir=tmp_path) - - def test_file_path_missing_file_part(self, tmp_path): - with pytest.raises(ConfigError, match="expected 'path/to/file.py:ClassName'"): - _load_runner_class(":MyRunner", base_dir=tmp_path) - - def test_file_not_found(self, tmp_path): - with pytest.raises(ConfigError, match="Runner file not found"): - _load_runner_class("does_not_exist.py:MyRunner", base_dir=tmp_path) - - def test_file_import_error(self, tmp_path): - bad = tmp_path / "bad_runner.py" - bad.write_text("raise RuntimeError('boom')\n") - with pytest.raises(ConfigError, match="Failed to import runner file"): - _load_runner_class("bad_runner.py:MyRunner", base_dir=tmp_path) - - def test_file_spec_none(self, tmp_path): - self._write_runner(tmp_path) +class TestMain: + def test_runs_yaml_and_prints_summary(self, tmp_path, capsys): + summary = RunSummary( + final_status=FinalStatus.PASSED, + iterations_run=2, + elapsed_s=1.25, + ) + yaml_path = tmp_path / "task.yaml" with patch( - "microbots.auto_memory.cli.importlib.util.spec_from_file_location", - return_value=None, - ): - with pytest.raises(ConfigError, match="Cannot load runner module"): - _load_runner_class("myrunner.py:MyRunner", base_dir=tmp_path) - - def test_file_class_not_found(self, tmp_path): - self._write_runner(tmp_path) - with pytest.raises(ConfigError, match="not found"): - _load_runner_class("myrunner.py:NoSuchRunner", base_dir=tmp_path) - - # --- dotted import path form --------------------------------------- - def test_dotted_form_loads_class(self, tmp_path): - cls = _load_runner_class( - "microbots.auto_memory.runners.writing_bot_runner.WritingBotRunner", - base_dir=tmp_path, + "microbots.auto_memory.cli.run_from_yaml", return_value=summary + ) as run: + assert main([str(yaml_path)]) == 0 + + run.assert_called_once_with( + yaml_path, + workdir=None, + run_id=None, + model=None, + external_memory_dir=None, ) - assert cls.__name__ == "WritingBotRunner" - - def test_dotted_form_missing_module(self, tmp_path): - with pytest.raises(ConfigError, match="expected"): - _load_runner_class("WritingBotRunner", base_dir=tmp_path) + assert "auto-memory passed: iterations=2 elapsed=1.2s" in capsys.readouterr().out - def test_dotted_form_import_error(self, tmp_path): - with pytest.raises(ConfigError, match="Cannot import runner module"): - _load_runner_class("no_such_pkg.module.Klass", base_dir=tmp_path) - - def test_dotted_form_import_time_error(self, tmp_path): - """A non-ImportError raised while importing the module is wrapped in - ConfigError instead of escaping as a raw traceback.""" + def test_overrides_yaml_options(self, tmp_path): + summary = RunSummary( + final_status=FinalStatus.PASSED, + iterations_run=1, + ) + yaml_path = tmp_path / "task.yaml" with patch( - "microbots.auto_memory.cli.importlib.import_module", - side_effect=RuntimeError("boom at import"), + "microbots.auto_memory.cli.run_from_yaml", return_value=summary + ) as run: + assert main([ + str(yaml_path), + "--model", _MODEL, + "--workdir", "work", + "--run-id", "fixed", + "--external-memory-dir", "memory", + ]) == 0 + + assert run.call_args.kwargs == { + "workdir": Path("work"), + "run_id": "fixed", + "model": _MODEL, + "external_memory_dir": Path("memory"), + } + + def test_reports_config_error(self, tmp_path, capsys): + with patch( + "microbots.auto_memory.cli.run_from_yaml", + side_effect=ConfigError("bad task"), ): - with pytest.raises(ConfigError, match="Failed to import runner module"): - _load_runner_class( - "microbots.auto_memory.config.TaskConfig", base_dir=tmp_path - ) - - def test_dotted_form_class_not_found(self, tmp_path): - with pytest.raises(ConfigError, match="not found"): - _load_runner_class( - "microbots.auto_memory.config.NoSuchClass", base_dir=tmp_path - ) + assert main([str(tmp_path / "task.yaml")]) == 2 + assert "config error: bad task" in capsys.readouterr().err + + def test_prints_summary_error(self, tmp_path, capsys): + summary = RunSummary( + final_status=FinalStatus.ERROR, + iterations_run=1, + error_message="runner failed", + ) + with patch("microbots.auto_memory.cli.run_from_yaml", return_value=summary): + assert main([str(tmp_path / "task.yaml")]) == 0 + assert "last error: runner failed" in capsys.readouterr().err + + def test_module_entry_point_dispatches_to_main(self): + with ( + patch("microbots.auto_memory.cli.main", return_value=7), + pytest.raises(SystemExit, match="7"), + ): + runpy.run_module("microbots.auto_memory.__main__", run_name="__main__") - def test_resolved_attribute_not_callable(self, tmp_path): - """A resolved attribute that is not callable (e.g. a constant) raises a - clear ConfigError instead of a later cryptic TypeError.""" - const_file = tmp_path / "const_runner.py" - const_file.write_text("NOT_A_RUNNER = 42\n") - with pytest.raises(ConfigError, match="not callable"): - _load_runner_class("const_runner.py:NOT_A_RUNNER", base_dir=tmp_path) diff --git a/test/auto_memory/test_config.py b/test/auto_memory/test_config.py index cf6e80ff..0e7ea62f 100644 --- a/test/auto_memory/test_config.py +++ b/test/auto_memory/test_config.py @@ -2,13 +2,12 @@ import pytest from pathlib import Path -from microbots.auto_memory.config import TaskConfig +from microbots.auto_memory.config import DEFAULT_PROMPT_TEMPLATE, TaskConfig from microbots.auto_memory.errors import ConfigError MINIMAL_YAML = textwrap.dedent("""\ task_definition: Fix the bug - prompt_template: "Goal: {{ task }}" callbacks: - name: tests command: pytest "$CANDIDATE" @@ -58,7 +57,7 @@ class TestLoadFromYaml: def test_minimal_valid(self, tmp_yaml): cfg = TaskConfig.load_from_yaml(tmp_yaml(MINIMAL_YAML)) assert cfg.task_definition == "Fix the bug" - assert "{{ task }}" in cfg.prompt_template + assert cfg.prompt_template == DEFAULT_PROMPT_TEMPLATE def test_full_yaml(self, tmp_yaml): cfg = TaskConfig.load_from_yaml(tmp_yaml(FULL_YAML)) @@ -86,6 +85,8 @@ def test_callbacks_parsed(self, tmp_yaml): def test_defaults_applied(self, tmp_yaml): cfg = TaskConfig.load_from_yaml(tmp_yaml(MINIMAL_YAML)) + assert cfg.model is None + assert cfg.workdir == ".auto-memory" assert cfg.max_iterations == 5 assert cfg.timeout_min == 60 assert cfg.per_iteration_timeout == 600 @@ -93,60 +94,19 @@ def test_defaults_applied(self, tmp_yaml): assert cfg.output_path == "candidate" assert cfg.reference_inputs == [] - def test_runner_defaults_to_writing_bot(self, tmp_yaml): - cfg = TaskConfig.load_from_yaml(tmp_yaml(MINIMAL_YAML)) - assert cfg.runner == ( - "microbots.auto_memory.runners.writing_bot_runner.WritingBotRunner" - ) - assert cfg.runner_params == {} - - def test_runner_fields_parsed(self, tmp_yaml): - yaml = textwrap.dedent("""\ - task_definition: Backport the patch - prompt_template: "Goal: {{ task }}" - runner: ./backport_runner.py:BackportRunner - runner_params: - repo_url: https://example.com/repo.git - target_commit: abc123 - callbacks: - - name: tests - command: pytest - """) + def test_model_and_workdir_parsed(self, tmp_yaml): + yaml = "model: azure-openai/gpt-4o\nworkdir: results\n" + MINIMAL_YAML cfg = TaskConfig.load_from_yaml(tmp_yaml(yaml)) - assert cfg.runner == "./backport_runner.py:BackportRunner" - assert cfg.runner_params == { - "repo_url": "https://example.com/repo.git", - "target_commit": "abc123", - } - - def test_runner_params_non_mapping_raises_config_error(self, tmp_yaml): - """A non-mapping runner_params in YAML surfaces as ConfigError, not a - raw ValueError/TypeError from dict() coercion.""" - yaml = textwrap.dedent("""\ - task_definition: Backport the patch - prompt_template: "Goal: {{ task }}" - runner_params: - - not - - a - - mapping - callbacks: - - name: tests - command: pytest - """) - with pytest.raises(ConfigError, match="runner_params.*mapping"): - TaskConfig.load_from_yaml(tmp_yaml(yaml)) + assert cfg.model == "azure-openai/gpt-4o" + assert cfg.workdir == "results" - def test_runner_params_scalar_raises_config_error(self, tmp_yaml): - yaml = textwrap.dedent("""\ - task_definition: Backport the patch - prompt_template: "Goal: {{ task }}" - runner_params: 5 - callbacks: - - name: tests - command: pytest - """) - with pytest.raises(ConfigError, match="runner_params.*mapping"): - TaskConfig.load_from_yaml(tmp_yaml(yaml)) + @pytest.mark.parametrize("field", ["runner", "runner_params"]) + def test_yaml_runner_fields_direct_users_to_python_injection( + self, tmp_yaml, field + ): + yaml = MINIMAL_YAML + f"{field}: custom\n" + with pytest.raises(ConfigError, match="agent_runner"): + TaskConfig.load_from_yaml(tmp_yaml(yaml)) def test_missing_callbacks(self, tmp_yaml): yaml = textwrap.dedent("""\ @@ -168,9 +128,10 @@ def test_missing_task_definition(self, tmp_yaml): with pytest.raises(ConfigError, match="task_definition"): TaskConfig.load_from_yaml(tmp_yaml("prompt_template: hello")) - def test_missing_prompt_template(self, tmp_yaml): - with pytest.raises(ConfigError, match="prompt_template"): - TaskConfig.load_from_yaml(tmp_yaml("task_definition: hello")) + def test_prompt_template_is_optional(self, tmp_yaml): + yaml = "task_definition: hello\ncallbacks:\n - name: check\n command: echo ok\n" + cfg = TaskConfig.load_from_yaml(tmp_yaml(yaml)) + assert cfg.prompt_template == DEFAULT_PROMPT_TEMPLATE def test_not_a_mapping(self, tmp_yaml): with pytest.raises(ConfigError, match="mapping"): @@ -202,6 +163,19 @@ def test_empty_prompt_template(self): with pytest.raises(ConfigError, match="prompt_template"): cfg.validate() + @pytest.mark.parametrize("model", ["gpt-4o", "unknown/gpt-4o"]) + def test_invalid_model(self, model): + cfg = self._base() + cfg.model = model + with pytest.raises(ConfigError, match="model"): + cfg.validate() + + def test_empty_workdir(self): + cfg = self._base() + cfg.workdir = "" + with pytest.raises(ConfigError, match="workdir"): + cfg.validate() + def test_max_iterations_zero(self): cfg = self._base() cfg.max_iterations = 0 @@ -309,20 +283,3 @@ def test_empty_callbacks_list(self): with pytest.raises(ConfigError, match="callbacks"): cfg.validate() - def test_runner_empty(self): - cfg = self._base() - cfg.runner = " " - with pytest.raises(ConfigError, match="runner"): - cfg.validate() - - def test_runner_params_not_a_dict(self): - cfg = self._base() - cfg.runner_params = ["not", "a", "dict"] - with pytest.raises(ConfigError, match="runner_params.*mapping"): - cfg.validate() - - def test_runner_params_reject_model_key(self): - cfg = self._base() - cfg.runner_params = {"model": "azure-openai/gpt-4o"} - with pytest.raises(ConfigError, match="runner_params.*model"): - cfg.validate() diff --git a/test/auto_memory/test_context.py b/test/auto_memory/test_context.py index 02dc3662..47d7f1c4 100644 --- a/test/auto_memory/test_context.py +++ b/test/auto_memory/test_context.py @@ -57,6 +57,29 @@ def test_unknown_variable_raises(self): @pytest.mark.unit class TestBuildIterationContextWithFeedback: + def test_default_template_renders_task_inputs_and_feedback(self): + cfg = TaskConfig( + task_definition="Fix the auth bug", + callbacks=[CallbackSpec(name="tests", command="pytest")], + reference_inputs=[ReferenceInput(name="spec", value="./spec.md")], + ) + feedback = Feedback( + iteration_idx=0, + summary="Authentication test failed", + root_causes=["Missing token validation"], + validator_failures=["test_auth.py failed"], + suggested_actions=["Validate the token before use"], + ) + + result = build_iteration_context(cfg, 1, feedback=feedback) + + assert "Fix the auth bug" in result + assert "spec: ./spec.md" in result + assert "Authentication test failed" in result + assert "Missing token validation" in result + assert "test_auth.py failed" in result + assert "Validate the token before use" in result + def test_feedback_summary_rendered(self): cfg = _config( "Goal: {{ task }}\n{% if feedback %}Feedback: {{ feedback.summary }}{% endif %}" diff --git a/test/auto_memory/test_loop.py b/test/auto_memory/test_loop.py new file mode 100644 index 00000000..0ae65a3e --- /dev/null +++ b/test/auto_memory/test_loop.py @@ -0,0 +1,145 @@ +"""Unit tests for the shared bounded iteration loop used by both orchestrators.""" + +from unittest.mock import patch + +import pytest + +from microbots.auto_memory.loop import StepOutcome, StopReason, run_bounded_loop + +pytestmark = pytest.mark.unit + + +class _Boom(Exception): + """Marker exception distinct from generic Exception for narrow-catch tests.""" + + +def test_limit_reached_when_no_step_ever_stops(): + calls = [] + + def step(idx): + calls.append(idx) + return StepOutcome(record=idx) + + result = run_bounded_loop(max_iterations=3, total_timeout_s=0, step=step) + + assert result.stop_reason == StopReason.LIMIT_REACHED + assert result.iterations_run == 3 + assert result.records == [0, 1, 2] + assert calls == [0, 1, 2] + + +def test_success_stop_ends_loop_immediately(): + def step(idx): + stop = StopReason.SUCCESS if idx == 1 else None + return StepOutcome(record=idx, stop=stop) + + result = run_bounded_loop(max_iterations=5, total_timeout_s=0, step=step) + + assert result.stop_reason == StopReason.SUCCESS + assert result.iterations_run == 2 + assert result.records == [0, 1] + + +def test_total_timeout_stops_before_running_next_iteration(): + def step(idx): + return StepOutcome(record=idx) + + with patch( + "microbots.auto_memory.loop.time.monotonic", + side_effect=[0.0, 0.0, 5.0, 61.0], + ): + result = run_bounded_loop(max_iterations=5, total_timeout_s=60, step=step) + + assert result.stop_reason == StopReason.TOTAL_TIMEOUT + # The iteration that would have exceeded the budget never ran. + assert result.iterations_run == 1 + assert result.records == [0] + + +def test_transient_error_retries_then_gives_up(): + attempts = [] + + def step(idx): + attempts.append(idx) + return StepOutcome(record=idx, transient_error=True) + + result = run_bounded_loop( + max_iterations=10, + total_timeout_s=0, + step=step, + max_transient_retries=2, + ) + + assert result.stop_reason == StopReason.ERROR + # 1 initial attempt + 2 retries = 3 iterations before giving up. + assert result.iterations_run == 3 + assert attempts == [0, 1, 2] + + +def test_transient_error_counter_resets_after_a_healthy_iteration(): + statuses = [True, False, True, True] # True == transient_error + + def step(idx): + return StepOutcome(record=idx, transient_error=statuses[idx]) + + result = run_bounded_loop( + max_iterations=len(statuses), + total_timeout_s=0, + step=step, + max_transient_retries=1, + ) + + # idx 0 transient (1/1, within budget) -> idx 1 healthy resets the + # counter -> idx 2 transient (1/1, within budget) -> idx 3 transient + # (2/1, exceeds budget) -> stop. + assert result.stop_reason == StopReason.ERROR + assert result.iterations_run == 4 + + +def test_caught_exception_produces_error_result_and_record(): + def step(idx): + raise _Boom("kaboom") + + def on_exception(idx, exc): + return f"iter {idx} failed: {exc}" + + result = run_bounded_loop( + max_iterations=5, + total_timeout_s=0, + step=step, + catch_exceptions=(_Boom,), + on_exception=on_exception, + ) + + assert result.stop_reason == StopReason.ERROR + assert result.iterations_run == 1 + assert result.records == ["iter 0 failed: kaboom"] + assert result.error_message == "kaboom" + + +def test_uncaught_exception_type_propagates(): + def step(idx): + raise ValueError("not caught") + + with pytest.raises(ValueError, match="not caught"): + run_bounded_loop( + max_iterations=5, + total_timeout_s=0, + step=step, + catch_exceptions=(_Boom,), + ) + + +def test_no_on_exception_callback_still_stops_without_a_record(): + def step(idx): + raise _Boom("kaboom") + + result = run_bounded_loop( + max_iterations=5, + total_timeout_s=0, + step=step, + catch_exceptions=(_Boom,), + ) + + assert result.stop_reason == StopReason.ERROR + assert result.records == [] diff --git a/test/auto_memory/test_memory.py b/test/auto_memory/test_memory.py index c9857c8b..7e9c511d 100644 --- a/test/auto_memory/test_memory.py +++ b/test/auto_memory/test_memory.py @@ -65,6 +65,74 @@ def test_mount_idempotent_called_twice(self, tmp_path): assert store.read_all() == [] +# --------------------------------------------------------------------------- +# external_memory_dir +# --------------------------------------------------------------------------- + +@pytest.mark.unit +class TestMemoryStoreExternalDir: + def test_external_dir_used_instead_of_run_subdir(self, tmp_path): + run_dir = tmp_path / "run" + ext = tmp_path / "trained_notes" + store = MemoryStore() + store.mount(run_dir, external_memory_dir=ext) + assert store.memory_dir == ext + assert ext.is_dir() + # The default sub-dir must NOT be created when external is used. + assert not (run_dir / "memory").exists() + + def test_external_dir_created_if_missing(self, tmp_path): + ext = tmp_path / "nested" / "trained" + store = MemoryStore() + store.mount(tmp_path / "run", external_memory_dir=ext) + assert ext.is_dir() + + def test_external_dir_preserves_pre_existing_notes(self, tmp_path): + ext = tmp_path / "trained" + ext.mkdir() + note = ext / "architecture.md" + note.write_text("# repo notes\n", encoding="utf-8") + + store = MemoryStore() + store.mount(tmp_path / "run", external_memory_dir=ext) + # Non-feedback files must be untouched. + assert note.read_text(encoding="utf-8") == "# repo notes\n" + + def test_external_dir_non_resume_wipes_only_feedback(self, tmp_path): + ext = tmp_path / "trained" + ext.mkdir() + note = ext / "notes.md" + note.write_text("keep me", encoding="utf-8") + stale_feedback = ext / "feedback.jsonl" + stale_feedback.write_text('{"iteration_idx": 9, "summary": "old"}\n', encoding="utf-8") + + store = MemoryStore() + store.mount(tmp_path / "run", resume=False, external_memory_dir=ext) + # Feedback is wiped ... + assert store.read_all() == [] + # ... but user's notes survive. + assert note.read_text(encoding="utf-8") == "keep me" + + def test_external_dir_resume_preserves_feedback(self, tmp_path): + ext = tmp_path / "trained" + store = MemoryStore() + store.mount(tmp_path / "run", external_memory_dir=ext) + store.append_feedback(_make_feedback(0, "prior")) + + store2 = MemoryStore() + store2.mount(tmp_path / "run2", resume=True, external_memory_dir=ext) + entries = store2.read_all() + assert len(entries) == 1 + assert entries[0].summary == "prior" + + def test_external_dir_accepts_string_path(self, tmp_path): + # Path-like objects: passing a string must also work. + ext = tmp_path / "trained" + store = MemoryStore() + store.mount(tmp_path / "run", external_memory_dir=str(ext)) # type: ignore[arg-type] + assert store.memory_dir == ext + + # --------------------------------------------------------------------------- # Unmounted guard # --------------------------------------------------------------------------- diff --git a/test/auto_memory/test_orchestrator.py b/test/auto_memory/test_orchestrator.py index d3db31b1..ff320f30 100644 --- a/test/auto_memory/test_orchestrator.py +++ b/test/auto_memory/test_orchestrator.py @@ -11,6 +11,7 @@ from microbots.auto_memory.config import TaskConfig from microbots.auto_memory.data_models import ( CallbackSpec, + Feedback, FinalStatus, IterationStatus, ) @@ -25,6 +26,22 @@ # --------------------------------------------------------------------------- +@pytest.fixture(autouse=True) +def _mock_failure_analysis(): + def _feedback(*, iteration_idx, **kwargs): + return Feedback( + iteration_idx=iteration_idx, + summary="Mock callback failure analysis", + validator_failures=["check"], + ) + + with patch( + "microbots.auto_memory.orchestrator.analyze_failure", + side_effect=_feedback, + ): + yield + + def _make_config( *, max_iterations: int = 5, @@ -66,7 +83,7 @@ def _make_callback_result(tmp_path: Path, *, passed: bool = True) -> CallbackRes ) -class MockAgentRunner: +class MockAgentRunner(AgentRunner): """Controllable AgentRunner that returns results from a pre-configured queue.""" def __init__(self, results: list[AgentResult]) -> None: @@ -84,7 +101,7 @@ def call_count(self) -> int: return len(self._calls) -class RaisingAgentRunner: +class RaisingAgentRunner(AgentRunner): """AgentRunner that always raises AgentError.""" def __init__(self, message: str = "agent exploded") -> None: @@ -339,7 +356,7 @@ def fake_monotonic(): # First call (start_time) returns 0; next call returns timeout + 1 return 0.0 if call_count == 1 else config.timeout_min * 60 + 1.0 - with patch("microbots.auto_memory.orchestrator.time.monotonic", side_effect=fake_monotonic): + with patch("microbots.auto_memory.loop.time.monotonic", side_effect=fake_monotonic): summary = orch.run() assert summary.final_status == FinalStatus.TIMEOUT diff --git a/test/auto_memory/test_workspace.py b/test/auto_memory/test_workspace.py index 1a653701..cef043f2 100644 --- a/test/auto_memory/test_workspace.py +++ b/test/auto_memory/test_workspace.py @@ -136,6 +136,50 @@ def test_prepare_resume_large_index(self, tmp_path): assert wm2.iteration_count == 101 +# --------------------------------------------------------------------------- +# external_memory_dir +# --------------------------------------------------------------------------- + +@pytest.mark.unit +class TestWorkspaceManagerExternalMemoryDir: + def test_external_dir_is_used_instead_of_run_subdir(self, tmp_path): + run_dir = tmp_path / "run" + ext = tmp_path / "trained" + wm = WorkspaceManager(run_dir=run_dir, external_memory_dir=ext) + wm.prepare() + assert wm.memory.memory_dir == ext + # The default sub-dir must not be created. + assert not (run_dir / "memory").exists() + + def test_external_dir_preserves_pre_existing_notes(self, tmp_path): + ext = tmp_path / "trained" + ext.mkdir() + (ext / "architecture.md").write_text("keep", encoding="utf-8") + + wm = WorkspaceManager(run_dir=tmp_path / "run", external_memory_dir=ext) + wm.prepare() + assert (ext / "architecture.md").read_text(encoding="utf-8") == "keep" + + def test_external_dir_survives_run_dir_wipe_on_non_resume(self, tmp_path): + """A second run with the same external memory keeps the notes intact.""" + ext = tmp_path / "trained" + ext.mkdir() + (ext / "notes.md").write_text("shared knowledge", encoding="utf-8") + + wm1 = WorkspaceManager(run_dir=tmp_path / "run1", external_memory_dir=ext) + wm1.prepare() + wm2 = WorkspaceManager(run_dir=tmp_path / "run2", external_memory_dir=ext) + wm2.prepare(resume=False) # would wipe internal memory; must not touch ext + assert (ext / "notes.md").read_text(encoding="utf-8") == "shared knowledge" + + def test_external_dir_default_is_none(self, tmp_path): + wm = WorkspaceManager(run_dir=tmp_path / "run") + assert wm.external_memory_dir is None + wm.prepare() + # Falls back to the traditional run_dir/memory/ layout. + assert wm.memory.memory_dir == tmp_path / "run" / "memory" + + # --------------------------------------------------------------------------- # prepare_iteration() / iteration_dir() # --------------------------------------------------------------------------- diff --git a/test/swe-bench-test/run_plain_baseline.py b/test/swe-bench-test/run_plain_baseline.py new file mode 100644 index 00000000..58be24bc --- /dev/null +++ b/test/swe-bench-test/run_plain_baseline.py @@ -0,0 +1,364 @@ +"""Minimal, single-shot plain-agent baseline for a SWE-bench repo. + +No orchestrator, no analyzer, no memory tool, no retries — this runs +``WritingBot`` exactly once per instance and scores the result using the +**official SWE-bench Docker evaluation harness** +(``swebench.harness.run_evaluation``, ``pip install swebench``). Docker +builds a per-instance image with the correct Python interpreter + deps for +that historical commit, so old instances (e.g. pre-3.8 code using the +removed ``imp`` module) are scored correctly without touching the host +Python at all — the host only needs Docker and the ``swebench`` pip package +(installed in this project's ``.venv``, never system-wide). + +Usage +----- + + python run_plain_baseline.py \\ + --repo pytest-dev/pytest \\ + --model azure-openai/gpt-5 \\ + --max-bot-steps 40 \\ + --timeout-s 3600 \\ + [--limit 5] [--instance-id pytest-dev__pytest-5262] + +Writes ``/results.json`` with one row per instance: +``{instance_id, resolved, bot_status, error}`` and prints a final +``resolved/total`` summary. Also leaves ``/predictions.jsonl`` and +the raw swebench harness report (``/..json``) for +inspection. +""" + +from __future__ import annotations + +import argparse +import json +import logging +import os +import subprocess +import sys +import time +from dataclasses import asdict, dataclass +from pathlib import Path + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../src"))) + +from microbots.bot.WritingBot import WritingBot + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", +) +logger = logging.getLogger(__name__) + +_TASK_TEMPLATE = """Fix the bug described below in the checked-out repository. +Your working directory holds a fresh clone at the buggy commit. + +## Problem statement + +{problem_statement} + +{hints_block} +## Guidance + +Produce a minimal patch that makes the target tests pass. Do NOT modify +files under `tests/` or `test/` unless the problem statement explicitly +requires it. Leave the working tree ready to be tested — do not commit. +""" + + +@dataclass +class InstanceResult: + instance_id: str + bot_status: bool + bot_error: str | None + resolved: bool + elapsed_s: float + + +# --------------------------------------------------------------------------- +# Repo setup +# --------------------------------------------------------------------------- + +def clone_and_checkout(repo: str, base_commit: str, dest: Path) -> None: + if dest.exists(): + _force_remove_dir(dest) + dest.parent.mkdir(parents=True, exist_ok=True) + subprocess.run( + ["git", "clone", "--quiet", f"https://github.com/{repo}.git", str(dest)], + check=True, + ) + subprocess.run( + ["git", "-C", str(dest), "checkout", "--quiet", base_commit], check=True + ) + + +def _force_remove_dir(dest: Path) -> None: + """Remove *dest*, tolerating root-owned leftovers from a prior Docker run. + + WritingBot mounts the repo dir into a container that runs as root, so any + files the agent's process creates inside (e.g. __pycache__/*.pyc) land on + the host owned by root and can't be removed by the current user. Falling + back to deleting via a throwaway container (root only inside that single + bind-mounted folder, nothing else on the host) avoids needing sudo. + """ + result = subprocess.run(["rm", "-rf", str(dest)]) + if result.returncode == 0: + return + logger.warning( + "Host rm -rf failed on %s (likely root-owned Docker leftovers); " + "retrying via a throwaway container.", dest, + ) + subprocess.run( + [ + "docker", "run", "--rm", + "-v", f"{dest}:/target", + "alpine", "sh", "-c", "rm -rf /target/* /target/..?* /target/.[!.]* 2>/dev/null; true", + ], + check=True, + ) + subprocess.run(["rm", "-rf", str(dest)], check=True) + + + +# --------------------------------------------------------------------------- +# Docker-based scoring via the official swebench harness +# --------------------------------------------------------------------------- + +_MODEL_NAME_OR_PATH = "microbots-plain-baseline" + + +def get_model_patch(repo_dir: Path) -> str: + """Return the bot's edits as a unified diff (empty string if none).""" + result = subprocess.run( + ["git", "-C", str(repo_dir), "diff"], + capture_output=True, text=True, check=True, + ) + return result.stdout + + +def run_swebench_harness( + predictions_path: Path, + dataset: str, + split: str, + instance_ids: list[str], + run_id: str, + out_dir: Path, + timeout_s: int, + max_workers: int = 1, +) -> dict: + """Invoke the official Docker-based swebench harness and return its report dict. + + Builds a per-instance Docker image with the correct Python interpreter + + deps for that historical commit, applies the model patch + gold + test_patch inside the container, and runs FAIL_TO_PASS/PASS_TO_PASS + there — so this is immune to host Python version mismatches (e.g. the + `imp` module removed in 3.12). The harness writes its report JSON into + the current working directory, so we run it with cwd=out_dir. + ``max_workers`` controls how many instance containers the harness runs + concurrently — safe to raise since each instance is fully isolated in + its own container. + """ + cmd = [ + sys.executable, "-m", "swebench.harness.run_evaluation", + "--dataset_name", dataset, + "--split", split, + "--predictions_path", str(predictions_path), + "--max_workers", str(max_workers), + "--run_id", run_id, + "--timeout", str(timeout_s), + "--instance_ids", *instance_ids, + ] + logger.info("Running swebench harness: %s", " ".join(cmd)) + subprocess.run(cmd, cwd=str(out_dir), check=True) + + report_path = out_dir / f"{_MODEL_NAME_OR_PATH}.{run_id}.json" + if not report_path.exists(): + raise FileNotFoundError( + f"swebench harness did not produce expected report at {report_path}" + ) + return json.loads(report_path.read_text(encoding="utf-8")) + + +def write_prediction(predictions_path: Path, instance_id: str, model_patch: str) -> None: + """Append/replace one instance's prediction in the shared predictions.jsonl.""" + existing: list[dict] = [] + if predictions_path.exists(): + existing = [json.loads(line) for line in predictions_path.read_text().splitlines() if line] + existing = [p for p in existing if p["instance_id"] != instance_id] + existing.append({ + "instance_id": instance_id, + "model_name_or_path": _MODEL_NAME_OR_PATH, + "model_patch": model_patch, + }) + predictions_path.write_text( + "\n".join(json.dumps(p) for p in existing) + "\n", encoding="utf-8" + ) + + +# --------------------------------------------------------------------------- +# Core +# --------------------------------------------------------------------------- + +def run_agent_once( + instance: dict, + *, + model: str, + work_root: Path, + max_bot_steps: int, + timeout_s: int, +) -> tuple[Path, "bool | None", str | None, float]: + """Clone + checkout + run WritingBot once. Returns + (repo_dir, bot_status, bot_error, elapsed_s). Scoring happens separately. + """ + instance_id = instance["instance_id"] + repo_dir = work_root / instance_id + started = time.monotonic() + + logger.info("=== %s: cloning + checkout (base_commit only, no test_patch) ===", instance_id) + clone_and_checkout(instance["repo"], instance["base_commit"], repo_dir) + + hints = (instance.get("hints_text") or "").strip() + hints_block = f"## Hints from maintainers\n\n{hints}\n\n" if hints else "" + task = _TASK_TEMPLATE.format( + problem_statement=instance["problem_statement"], + hints_block=hints_block, + ) + + logger.info("=== %s: running plain agent (1 shot, max_bot_steps=%d) ===", + instance_id, max_bot_steps) + bot = WritingBot(model=model, folder_to_mount=str(repo_dir)) + try: + bot_result = bot.run(task=task, max_iterations=max_bot_steps, timeout_in_seconds=timeout_s) + finally: + if bot.environment is not None: + bot.environment.stop() + + elapsed = time.monotonic() - started + return repo_dir, bot_result.status, bot_result.error, elapsed + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + +def _build_arg_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--dataset", default="princeton-nlp/SWE-bench_Verified") + p.add_argument("--split", default="test") + p.add_argument("--repo", default=None, + help='e.g. "pytest-dev/pytest". Required unless --instance-id is given.') + p.add_argument("--instance-id", default=None, help="Run just this one instance.") + p.add_argument("--limit", type=int, default=None) + p.add_argument("--model", required=True, help="e.g. azure-openai/gpt-5") + p.add_argument("--max-bot-steps", type=int, default=40, + help="WritingBot's own step budget for the single attempt.") + p.add_argument("--timeout-s", type=int, default=3600, + help="Both the bot's wall-clock budget and the harness's per-test timeout.") + p.add_argument("--work-dir", type=Path, default=Path("/tmp/plain_baseline")) + p.add_argument("--out-dir", type=Path, default=Path("/tmp/plain_baseline_results")) + p.add_argument("--run-id", default=None, + help="swebench harness run_id (default: 'plain-baseline-').") + p.add_argument("--max-workers", type=int, default=1, + help="Concurrent Docker containers for the scoring phase (default: 1). " + "Each instance is fully isolated, so raising this is safe as long " + "as your machine has the CPU/RAM/disk to build+run that many images " + "at once.") + return p + + +def main(argv: list[str] | None = None) -> int: + args = _build_arg_parser().parse_args(argv) + if not args.instance_id and not args.repo: + print("Must pass --repo or --instance-id", file=sys.stderr) + return 2 + + from datasets import load_dataset + + ds = load_dataset(args.dataset, split=args.split) + if args.instance_id: + instances = [dict(r) for r in ds if r["instance_id"] == args.instance_id] + else: + instances = [dict(r) for r in ds if r["repo"] == args.repo] + if args.limit is not None: + instances = instances[: args.limit] + + if not instances: + print(f"No instances found for repo={args.repo!r}", file=sys.stderr) + return 1 + + args.work_dir.mkdir(parents=True, exist_ok=True) + args.out_dir.mkdir(parents=True, exist_ok=True) + predictions_path = args.out_dir / "predictions.jsonl" + run_id = args.run_id or f"plain-baseline-{int(time.time())}" + + # --- Phase 1: run the plain agent once per instance, collect patches --- + bot_meta: dict[str, dict] = {} + for i, inst in enumerate(instances, start=1): + instance_id = inst["instance_id"] + logger.info("--- agent %d/%d: %s ---", i, len(instances), instance_id) + try: + repo_dir, bot_status, bot_error, elapsed = run_agent_once( + inst, + model=args.model, + work_root=args.work_dir, + max_bot_steps=args.max_bot_steps, + timeout_s=args.timeout_s, + ) + model_patch = get_model_patch(repo_dir) if bot_status else "" + except Exception as exc: # noqa: BLE001 — one bad instance shouldn't kill the batch + logger.exception("%s: unhandled error during agent run", instance_id) + bot_status, bot_error, elapsed, model_patch = False, f"{type(exc).__name__}: {exc}", 0.0, "" + + bot_meta[instance_id] = { + "bot_status": bool(bot_status), + "bot_error": bot_error, + "elapsed_s": elapsed, + } + write_prediction(predictions_path, instance_id, model_patch) + + # --- Phase 2: score every instance in one Docker harness invocation --- + instance_ids = [inst["instance_id"] for inst in instances] + logger.info("=== scoring %d instance(s) via swebench Docker harness ===", len(instance_ids)) + try: + report = run_swebench_harness( + predictions_path, args.dataset, args.split, instance_ids, + run_id, args.out_dir, args.timeout_s, max_workers=args.max_workers, + ) + except Exception: + logger.exception("swebench harness invocation failed") + report = {"resolved_ids": [], "error_ids": instance_ids} + + resolved_ids = set(report.get("resolved_ids", [])) + error_ids = set(report.get("error_ids", [])) + + # --- Combine + report --- + results: list[InstanceResult] = [] + for inst in instances: + instance_id = inst["instance_id"] + meta = bot_meta.get(instance_id, {}) + resolved = instance_id in resolved_ids + results.append(InstanceResult( + instance_id=instance_id, + bot_status=meta.get("bot_status", False), + bot_error=meta.get("bot_error") or ("harness_error" if instance_id in error_ids else None), + resolved=resolved, + elapsed_s=meta.get("elapsed_s", 0.0), + )) + + (args.out_dir / "results.json").write_text( + json.dumps([asdict(r) for r in results], indent=2), encoding="utf-8" + ) + + resolved_count = sum(1 for r in results if r.resolved) + print() + print(f"{'instance_id':45s} resolved bot_ok elapsed") + print("-" * 80) + for r in results: + print(f"{r.instance_id:45s} {str(r.resolved):8s} {str(r.bot_status):6s} {r.elapsed_s:6.0f}s") + print("-" * 80) + print(f"resolved: {resolved_count}/{len(results)}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) +