tinymind-native-8b-remote-handoff/bundle/evaluation/axiom_orchestrator.py

"""TinyMind Axiom Orchestrator.

An original, deterministic orchestration runtime that joins workflow planning,
safety gating, sandbox execution, evidence ledgers, and multi-agent critique
into one auditable harness. It is intentionally small and local-first: the goal
is a reliable training/eval substrate, not an unbounded autonomous host agent.
"""

from __future__ import annotations

from collections import Counter
from dataclasses import asdict, dataclass
from datetime import datetime, timezone
import hashlib
import json
from pathlib import Path
import tempfile
from typing import Any

from model.sandbox_tool_core import SandboxToolCore, SandboxToolPolicy


SCHEMA_VERSION = "tinymind-axiom-orchestrator-v1"


def _now() -> str:
    return datetime.now(timezone.utc).isoformat()


def _sha(payload: Any) -> str:
    return hashlib.sha256(json.dumps(payload, ensure_ascii=False, sort_keys=True).encode("utf-8")).hexdigest()


@dataclass(frozen=True)
class AxiomAgent:
    name: str
    role: str
    objective: str
    veto_domains: tuple[str, ...] = ()

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


@dataclass(frozen=True)
class AxiomStep:
    step_id: str
    kind: str
    owner: str
    goal: str
    tool: str | None = None
    args: dict[str, Any] | None = None
    requires_evidence: bool = True

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


DEFAULT_AGENTS = (
    AxiomAgent("architect", "workflow designer", "decompose tasks into minimal verifiable steps"),
    AxiomAgent("executor", "sandbox operator", "execute only policy-approved local actions"),
    AxiomAgent("critic", "failure analyst", "find missing evidence, unsafe actions, and weak claims"),
    AxiomAgent("evidence", "ledger keeper", "bind every outcome to hashes and exact artifacts"),
)


class AxiomOrchestrator:
    """One-pass deterministic orchestration harness for TinyMind workflows."""

    def __init__(self, root: str | Path, *, agents: tuple[AxiomAgent, ...] = DEFAULT_AGENTS):
        self.root = Path(root).resolve()
        self.root.mkdir(parents=True, exist_ok=True)
        self.agents = agents
        self.sandbox_root = self.root / "sandbox"
        self.core = SandboxToolCore(
            self.sandbox_root,
            policy=SandboxToolPolicy(
                command_allowlist=("echo",),
                max_write_bytes=64_000,
                max_files_per_project=24,
                cmd_timeout_s=3.0,
                allow_public_internet=False,
            ),
        )

    def plan(self, mission: str) -> list[AxiomStep]:
        safe_name = _safe_name(mission)
        return [
            AxiomStep("S1", "contract", "architect", "write mission contract and success checks"),
            AxiomStep(
                "S2",
                "sandbox",
                "executor",
                "create isolated workspace for the mission",
                "sandbox.env.create",
                {"name": safe_name},
            ),
            AxiomStep(
                "S3",
                "artifact",
                "executor",
                "write a minimal evidence artifact inside the sandbox",
                "sandbox.env.file_put",
                {"name": safe_name, "path": "MISSION.md", "content": _mission_doc(mission)},
            ),
            AxiomStep(
                "S4",
                "harness",
                "executor",
                "run deterministic smoke harness",
                "sandbox.env.run",
                {"name": safe_name, "argv": ["echo", "axiom-harness-ok"]},
            ),
            AxiomStep(
                "S5",
                "snapshot",
                "evidence",
                "snapshot the verified sandbox state",
                "sandbox.env.snapshot",
                {"name": safe_name, "snapshot": "verified"},
            ),
            AxiomStep("S6", "critique", "critic", "score evidence and block unsupported claims"),
        ]

    def run(self, mission: str) -> dict[str, Any]:
        steps = self.plan(mission)
        events: list[dict[str, Any]] = []
        critiques: list[dict[str, Any]] = []
        for step in steps:
            before = {"step": step.to_dict(), "mission_sha256": _sha(mission), "event": "before"}
            if step.tool:
                result = self.core.call(step.tool, step.args or {})
            else:
                result = self._internal_step(step, mission, events)
            event = {
                "step": step.to_dict(),
                "ok": bool(result.get("ok", True)) if isinstance(result, dict) else True,
                "result": _compact_result(result),
                "input_sha256": _sha(before),
                "output_sha256": _sha(result),
            }
            events.append(event)
            critiques.append(self._critique(step, event))

        checks = self._checks(events, critiques)
        return {
            "schema_version": SCHEMA_VERSION,
            "created_at": _now(),
            "mission": mission,
            "agents": [agent.to_dict() for agent in self.agents],
            "plan": [step.to_dict() for step in steps],
            "events": events,
            "critiques": critiques,
            "checks": checks,
            "metrics": self._metrics(events, critiques),
            "sandbox_ledger_path": str(self.core.ledger_path),
            "claim_gate": {
                "axiom_orchestrator_ready": all(checks.values()),
                "autonomous_unbounded_claim_allowed": False,
                "world_best_orchestration_claim_allowed": False,
                "reason": "This runtime proves local deterministic orchestration with policy-gated tools; external comparison is still required.",
            },
        }

    def _internal_step(self, step: AxiomStep, mission: str, events: list[dict[str, Any]]) -> dict[str, Any]:
        if step.kind == "contract":
            return {
                "ok": True,
                "contract": {
                    "mission_sha256": _sha(mission),
                    "success_checks": [
                        "all steps emit hashes",
                        "sandbox tool policy blocks public internet by default",
                        "critic veto prevents unsupported claims",
                        "verified snapshot exists",
                    ],
                },
            }
        if step.kind == "critique":
            return {
                "ok": True,
                "observed_events": len(events),
                "missing_hashes": sum(1 for event in events if not event.get("input_sha256") or not event.get("output_sha256")),
            }
        return {"ok": False, "error": "unknown_internal_step"}

    def _critique(self, step: AxiomStep, event: dict[str, Any]) -> dict[str, Any]:
        problems: list[str] = []
        if step.requires_evidence and (not event.get("input_sha256") or not event.get("output_sha256")):
            problems.append("missing_hash_evidence")
        if not event.get("ok"):
            problems.append("step_failed")
        if step.tool == "sandbox.proxy":
            problems.append("network_proxy_requires_explicit_review")
        return {
            "step_id": step.step_id,
            "critic": "critic",
            "passed": not problems,
            "problems": problems,
            "veto": bool(problems),
        }

    def _checks(self, events: list[dict[str, Any]], critiques: list[dict[str, Any]]) -> dict[str, bool]:
        tool_manifest = self.core.manifest()
        event_kinds = {event["step"]["kind"] for event in events}
        return {
            "workflow_graph_present": {"contract", "sandbox", "artifact", "harness", "snapshot", "critique"} <= event_kinds,
            "all_steps_hashed": all(event.get("input_sha256") and event.get("output_sha256") for event in events),
            "sandbox_policy_gated": tool_manifest["policy"]["allow_public_internet"] is False,
            "sandbox_ledger_exists": self.core.ledger_path.exists(),
            "critic_veto_path_present": all("veto" in critique for critique in critiques),
            "all_critiques_passed": all(critique["passed"] for critique in critiques),
        }

    def _metrics(self, events: list[dict[str, Any]], critiques: list[dict[str, Any]]) -> dict[str, Any]:
        kinds = Counter(event["step"]["kind"] for event in events)
        passed = sum(1 for event in events if event.get("ok"))
        return {
            "steps_total": len(events),
            "steps_passed": passed,
            "step_pass_rate": round(passed / max(1, len(events)), 6),
            "agent_count": len(self.agents),
            "event_kinds": dict(sorted(kinds.items())),
            "critic_vetoes": sum(1 for critique in critiques if critique["veto"]),
            "collaboration_density": round(len(self.agents) / max(1, len(events)), 6),
        }


def build_axiom_orchestrator_report(
    out_dir: str | Path,
    *,
    mission: str = "Build a verified local TinyMind workflow with sandbox evidence.",
) -> dict[str, Any]:
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    runtime_root = Path(tempfile.mkdtemp(prefix="axiom_orchestrator_", dir=out))
    report = AxiomOrchestrator(runtime_root).run(mission)
    sft_path = out / "axiom_orchestrator_sft.jsonl"
    _write_sft(sft_path, report)
    report["sft_path"] = str(sft_path)
    json_path = out / "axiom_orchestrator_report.json"
    md_path = out / "axiom_orchestrator_report.md"
    report["json_path"] = str(json_path)
    report["markdown_path"] = str(md_path)
    json_path.write_text(json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True), encoding="utf-8")
    md_path.write_text(_markdown(report), encoding="utf-8")
    return report


def _safe_name(mission: str) -> str:
    digest = _sha(mission)[:10]
    return f"axiom-{digest}"


def _mission_doc(mission: str) -> str:
    return "\n".join(
        [
            "# Axiom Mission",
            "",
            mission.strip(),
            "",
            "Success requires deterministic tool output, safety gates, hash evidence, and critic approval.",
            "",
        ]
    )


def _compact_result(result: Any) -> Any:
    if not isinstance(result, dict):
        return result
    keep = {key: result[key] for key in ("ok", "error", "result", "stdout", "stderr", "input_sha256", "output_sha256") if key in result}
    if "result" in keep and isinstance(keep["result"], dict):
        keep["result"] = {
            key: keep["result"][key]
            for key in ("name", "path", "snapshot", "stdout", "stderr", "exit_code", "auto_snapshot")
            if key in keep["result"]
        } or {"keys": sorted(keep["result"].keys())[:12]}
    return keep or {"keys": sorted(result.keys())[:12]}


def _write_sft(path: Path, report: dict[str, Any]) -> None:
    row = {
        "messages": [
            {
                "role": "system",
                "content": "You are TinyMind Axiom Orchestrator. Plan, execute through sandbox tools, cite evidence hashes, and let critic veto weak claims.",
            },
            {"role": "user", "content": report["mission"]},
            {
                "role": "assistant",
                "content": json.dumps(
                    {
                        "plan": report["plan"],
                        "metrics": report["metrics"],
                        "checks": report["checks"],
                        "claim_gate": report["claim_gate"],
                    },
                    ensure_ascii=False,
                    indent=2,
                    sort_keys=True,
                ),
            },
        ],
        "source": "axiom_orchestrator",
        "metadata": {
            "domain": "sandbox_tools",
            "loss_weight": 1.35,
            "schema_version": report["schema_version"],
            "mission_sha256": _sha(report["mission"]),
        },
    }
    path.write_text(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n", encoding="utf-8")


def _markdown(report: dict[str, Any]) -> str:
    lines = [
        "# TinyMind Axiom Orchestrator",
        "",
        f"- Ready: {report['claim_gate']['axiom_orchestrator_ready']}",
        f"- Steps: {report['metrics']['steps_passed']}/{report['metrics']['steps_total']}",
        f"- Agents: {report['metrics']['agent_count']}",
        f"- Critic vetoes: {report['metrics']['critic_vetoes']}",
        f"- Sandbox ledger: `{report['sandbox_ledger_path']}`",
        f"- SFT: `{report['sft_path']}`",
        f"- World-best claim allowed: {report['claim_gate']['world_best_orchestration_claim_allowed']}",
        "",
        "## Checks",
        "",
    ]
    for key, value in report["checks"].items():
        lines.append(f"- {key}: {value}")
    return "\n".join(lines) + "\n"