L'électron rare
408 lines
15 KiB
Python
408 lines
15 KiB
Python
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import json
|
|
import re
|
|
import threading
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
from .state import load_fleet_config, load_models_config
|
|
|
|
|
|
def _fleet_labels() -> dict[str, str]:
|
|
fleet = load_fleet_config()
|
|
labels = [machine["label"] for machine in fleet["machines"]]
|
|
local_machine = fleet.get("local_machine", labels[0] if labels else "command-center")
|
|
return {
|
|
"local": local_machine,
|
|
"builder_a": labels[1] if len(labels) > 1 else "builder-a",
|
|
"builder_b": labels[2] if len(labels) > 2 else "builder-b",
|
|
"edge_a": labels[3] if len(labels) > 3 else "edge-a",
|
|
"lab_a": labels[4] if len(labels) > 4 else "lab-a",
|
|
}
|
|
|
|
|
|
def _strip_code_fences(text: str) -> str:
|
|
stripped = text.strip()
|
|
if stripped.startswith("```"):
|
|
stripped = re.sub(r"^```[a-zA-Z0-9_-]*\n", "", stripped)
|
|
stripped = re.sub(r"\n```$", "", stripped)
|
|
return stripped.strip()
|
|
|
|
|
|
def _extract_json(text: str) -> dict[str, Any]:
|
|
stripped = _strip_code_fences(text)
|
|
try:
|
|
return json.loads(stripped)
|
|
except json.JSONDecodeError:
|
|
start = stripped.find("{")
|
|
end = stripped.rfind("}")
|
|
if start >= 0 and end > start:
|
|
return json.loads(stripped[start : end + 1])
|
|
raise
|
|
|
|
|
|
def _lower(value: str) -> str:
|
|
return value.lower()
|
|
|
|
|
|
def _route_machine(goal: str) -> str:
|
|
labels = _fleet_labels()
|
|
goal_l = _lower(goal)
|
|
if any(token in goal_l for token in ("ssh", "shell", "fleet", "ops", "network", "infra")):
|
|
return labels["edge_a"]
|
|
if any(token in goal_l for token in ("refactor", "performance", "deep", "cuda", "compiler", "benchmark")):
|
|
return labels["builder_b"]
|
|
if any(token in goal_l for token in ("review", "audit", "test", "qa", "docs")):
|
|
return labels["builder_a"]
|
|
if any(token in goal_l for token in ("experiment", "prototype", "sandbox", "branch")):
|
|
return labels["lab_a"]
|
|
return labels["local"]
|
|
|
|
|
|
def _route_cli(goal: str) -> str:
|
|
goal_l = _lower(goal)
|
|
if any(token in goal_l for token in ("brainstorm", "idea", "alternative", "explore", "compare")):
|
|
return "vibe"
|
|
if any(token in goal_l for token in ("spec", "summary", "summarize", "review", "docs")):
|
|
return "claude"
|
|
return "codex"
|
|
|
|
|
|
def _route_backend(machine: str, cli: str, goal: str) -> str:
|
|
local_machine = _fleet_labels()["local"]
|
|
goal_l = _lower(goal)
|
|
if machine == local_machine and cli in {"claude", "codex"} and any(token in goal_l for token in ("parallel", "team", "multi-agent", "compare")):
|
|
return "omc"
|
|
if machine == local_machine:
|
|
return "cli-local"
|
|
return "cli-ssh"
|
|
|
|
|
|
def _route_agent(cli: str) -> str:
|
|
return {
|
|
"claude": "planner",
|
|
"codex": "executor",
|
|
"vibe": "analyst",
|
|
"copilot": "assistant",
|
|
}.get(cli, "executor")
|
|
|
|
|
|
def _heuristic_plan(goal: str, workspace: str, mission_id: str) -> dict[str, Any]:
|
|
machine = _route_machine(goal)
|
|
cli = _route_cli(goal)
|
|
backend = _route_backend(machine, cli, goal)
|
|
return {
|
|
"title": goal.strip().splitlines()[0][:80] or mission_id,
|
|
"summary": f"Route the mission to {machine} using {cli} via {backend}.",
|
|
"steps": [
|
|
{
|
|
"step_id": "step-1",
|
|
"kind": "task",
|
|
"title": f"Execute mission on {machine}",
|
|
"prompt": goal,
|
|
"machine": machine,
|
|
"backend": backend,
|
|
"cli": cli,
|
|
"agent": _route_agent(cli),
|
|
"status": "pending",
|
|
"attempt": 0,
|
|
"depends_on": [],
|
|
"llama_plan": "heuristic planner fallback",
|
|
"qwen_review": "",
|
|
"result_summary": "",
|
|
"artifact_refs": [],
|
|
}
|
|
],
|
|
}
|
|
|
|
|
|
def _heuristic_buddy_review(goal: str, plan: dict[str, Any]) -> dict[str, Any]:
|
|
step = plan["steps"][0]
|
|
cli = step["cli"]
|
|
machine = step["machine"]
|
|
local_machine = _fleet_labels()["local"]
|
|
suggestions: list[str] = []
|
|
agrees = True
|
|
|
|
if cli == "copilot":
|
|
agrees = False
|
|
suggestions.append("copilot is manual-only in v1; use codex or claude instead")
|
|
if any(token in _lower(goal) for token in ("fix", "implement", "refactor", "bug")) and cli != "codex":
|
|
agrees = False
|
|
suggestions.append("technical execution is better routed to codex")
|
|
if machine == local_machine and step["backend"] == "cli-ssh":
|
|
agrees = False
|
|
suggestions.append("local machine should use cli-local or omc, not cli-ssh")
|
|
|
|
summary = "Buddy review agrees with the plan." if agrees else "; ".join(suggestions)
|
|
return {
|
|
"agrees": agrees,
|
|
"summary": summary,
|
|
"suggested_cli": "codex" if any("codex" in entry for entry in suggestions) else None,
|
|
"suggested_backend": "cli-local" if machine == local_machine and step["backend"] == "cli-ssh" else None,
|
|
}
|
|
|
|
|
|
def _heuristic_verify(step: dict[str, Any], execution: dict[str, Any]) -> dict[str, Any]:
|
|
stdout = execution.get("stdout", "")
|
|
stderr = execution.get("stderr", "")
|
|
combined = f"{stdout}\n{stderr}".lower()
|
|
return_code = execution.get("returncode", 1)
|
|
|
|
if any(token in combined for token in ("not logged in", "/login", "device-auth", "device auth", "authentication required", "please login", "please log in")):
|
|
return {
|
|
"decision": "needs_human",
|
|
"summary": "The CLI needs an interactive login or credentials refresh.",
|
|
"confidence": "high",
|
|
}
|
|
|
|
if return_code != 0:
|
|
if step.get("attempt", 0) <= 1:
|
|
return {
|
|
"decision": "retry",
|
|
"summary": "Execution failed once; retry is reasonable.",
|
|
"confidence": "medium",
|
|
}
|
|
return {
|
|
"decision": "failed",
|
|
"summary": "Execution failed after retry budget was exhausted.",
|
|
"confidence": "high",
|
|
}
|
|
|
|
if "error" in combined and "success" not in combined:
|
|
return {
|
|
"decision": "needs_human",
|
|
"summary": "Output contains an error marker despite zero exit status.",
|
|
"confidence": "medium",
|
|
}
|
|
|
|
return {
|
|
"decision": "done",
|
|
"summary": "Execution completed successfully.",
|
|
"confidence": "medium",
|
|
}
|
|
|
|
|
|
def _heuristic_buddy_answer(prompt: str, mission: dict[str, Any] | None) -> dict[str, Any]:
|
|
title = mission["title"] if mission else "mission"
|
|
return {
|
|
"answer": f"Qwen buddy heuristic view for {title}: focus on route correctness, CLI fit, and whether codex should own the technical step. Prompt: {prompt}",
|
|
"mode": "heuristic",
|
|
}
|
|
|
|
|
|
@dataclass
|
|
class ModelHealth:
|
|
role: str
|
|
model_id: str
|
|
available: bool
|
|
loaded: bool
|
|
backend: str
|
|
|
|
|
|
class PlannerEngine:
|
|
def __init__(self) -> None:
|
|
self._models = load_models_config()
|
|
self._loaded: dict[str, tuple[Any, Any]] = {}
|
|
self._lock = threading.Lock()
|
|
self._mlx_python = importlib.util.find_spec("mlx_lm") is not None
|
|
|
|
def health(self) -> dict[str, Any]:
|
|
health = {}
|
|
for role in ("planner", "buddy", "verify"):
|
|
spec = self._models[role]
|
|
health[role] = ModelHealth(
|
|
role=role,
|
|
model_id=spec["model_id"],
|
|
available=self._mlx_python,
|
|
loaded=role in self._loaded,
|
|
backend="mlx-python" if self._mlx_python else "heuristic",
|
|
).__dict__
|
|
return {
|
|
"mlx_python": self._mlx_python,
|
|
"models": health,
|
|
"fallback_mode": self._models.get("fallback_mode", "heuristic"),
|
|
}
|
|
|
|
def _load_model(self, role: str) -> tuple[Any, Any]:
|
|
if role in self._loaded:
|
|
return self._loaded[role]
|
|
|
|
if not self._mlx_python:
|
|
raise RuntimeError("mlx_lm is not available")
|
|
|
|
with self._lock:
|
|
if role in self._loaded:
|
|
return self._loaded[role]
|
|
from mlx_lm import load # type: ignore
|
|
|
|
model_id = self._models[role]["model_id"]
|
|
model, tokenizer = load(model_id)
|
|
self._loaded[role] = (model, tokenizer)
|
|
return model, tokenizer
|
|
|
|
def _call_model_json(self, role: str, system_prompt: str, user_prompt: str) -> dict[str, Any]:
|
|
if not self._mlx_python:
|
|
raise RuntimeError("mlx_lm is not available")
|
|
|
|
model, tokenizer = self._load_model(role)
|
|
from mlx_lm import generate # type: ignore
|
|
|
|
messages = [
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": user_prompt},
|
|
]
|
|
prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
|
|
text = generate(
|
|
model,
|
|
tokenizer,
|
|
prompt=prompt,
|
|
max_tokens=self._models[role]["max_tokens"],
|
|
verbose=False,
|
|
)
|
|
return _extract_json(text)
|
|
|
|
def _call_model_text(self, role: str, system_prompt: str, user_prompt: str) -> str:
|
|
if not self._mlx_python:
|
|
raise RuntimeError("mlx_lm is not available")
|
|
|
|
model, tokenizer = self._load_model(role)
|
|
from mlx_lm import generate # type: ignore
|
|
|
|
messages = [
|
|
{"role": "system", "content": system_prompt},
|
|
{"role": "user", "content": user_prompt},
|
|
]
|
|
prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
|
|
return generate(
|
|
model,
|
|
tokenizer,
|
|
prompt=prompt,
|
|
max_tokens=self._models[role]["max_tokens"],
|
|
verbose=False,
|
|
).strip()
|
|
|
|
def plan_mission(self, mission: dict[str, Any]) -> dict[str, Any]:
|
|
plan = _heuristic_plan(mission["goal"], mission["workspace"], mission["mission_id"])
|
|
review = _heuristic_buddy_review(mission["goal"], plan)
|
|
fleet = load_fleet_config()
|
|
machine_labels = [machine["label"] for machine in fleet["machines"]]
|
|
local_machine = fleet.get("local_machine", machine_labels[0] if machine_labels else "command-center")
|
|
|
|
if self._mlx_python:
|
|
system = (
|
|
"You are Llama, the main orchestrator for a 5-machine coding fleet. "
|
|
"Return strict JSON with keys: title, summary, steps. "
|
|
"Each step must include: step_id, kind, title, prompt, machine, backend, cli, agent, depends_on."
|
|
)
|
|
user = json.dumps(
|
|
{
|
|
"mission_id": mission["mission_id"],
|
|
"goal": mission["goal"],
|
|
"workspace": mission["workspace"],
|
|
"machines": machine_labels,
|
|
"allowed_backends": ["omc", "cli-local", "cli-ssh", "zellij"],
|
|
"allowed_clis": ["claude", "codex", "vibe"],
|
|
},
|
|
indent=2,
|
|
)
|
|
try:
|
|
llama_plan = self._call_model_json("planner", system, user)
|
|
if "steps" in llama_plan and llama_plan["steps"]:
|
|
plan = {
|
|
"title": llama_plan.get("title", plan["title"]),
|
|
"summary": llama_plan.get("summary", plan["summary"]),
|
|
"steps": [],
|
|
}
|
|
for index, step in enumerate(llama_plan["steps"], start=1):
|
|
plan["steps"].append(
|
|
{
|
|
"step_id": step.get("step_id", f"step-{index}"),
|
|
"kind": step.get("kind", "task"),
|
|
"title": step.get("title", f"Step {index}"),
|
|
"prompt": step.get("prompt", mission["goal"]),
|
|
"machine": step.get("machine", _route_machine(mission["goal"])),
|
|
"backend": step.get("backend", _route_backend(step.get("machine", local_machine), step.get("cli", "codex"), mission["goal"])),
|
|
"cli": step.get("cli", _route_cli(mission["goal"])),
|
|
"agent": step.get("agent", _route_agent(step.get("cli", "codex"))),
|
|
"status": "pending",
|
|
"attempt": 0,
|
|
"depends_on": step.get("depends_on", []),
|
|
"llama_plan": llama_plan.get("summary", ""),
|
|
"qwen_review": "",
|
|
"result_summary": "",
|
|
"artifact_refs": [],
|
|
}
|
|
)
|
|
except Exception:
|
|
pass
|
|
|
|
buddy_system = (
|
|
"You are Qwen, the local technical buddy. Return strict JSON with keys: "
|
|
"agrees, summary, suggested_cli, suggested_backend."
|
|
)
|
|
buddy_user = json.dumps({"goal": mission["goal"], "plan": plan}, indent=2)
|
|
try:
|
|
review = self._call_model_json("buddy", buddy_system, buddy_user)
|
|
except Exception:
|
|
pass
|
|
|
|
for step in plan["steps"]:
|
|
step["qwen_review"] = review.get("summary", "")
|
|
if not review.get("agrees", True):
|
|
if review.get("suggested_cli") and step["cli"] == "claude":
|
|
step["cli"] = review["suggested_cli"]
|
|
step["agent"] = _route_agent(step["cli"])
|
|
if review.get("suggested_backend"):
|
|
step["backend"] = review["suggested_backend"]
|
|
|
|
return {
|
|
"plan": plan,
|
|
"buddy_review": review,
|
|
}
|
|
|
|
def verify_step(self, mission: dict[str, Any], step: dict[str, Any], execution: dict[str, Any]) -> dict[str, Any]:
|
|
result = _heuristic_verify(step, execution)
|
|
|
|
if self._mlx_python:
|
|
system = (
|
|
"You are the orchestrator verifier. Return strict JSON with keys: "
|
|
"decision, summary, confidence. Decisions: done, retry, failed, needs_human."
|
|
)
|
|
user = json.dumps(
|
|
{
|
|
"mission_title": mission["title"],
|
|
"step": step,
|
|
"execution": {
|
|
"returncode": execution.get("returncode"),
|
|
"stdout_tail": execution.get("stdout", "")[-4000:],
|
|
"stderr_tail": execution.get("stderr", "")[-4000:],
|
|
},
|
|
},
|
|
indent=2,
|
|
)
|
|
try:
|
|
result = self._call_model_json("verify", system, user)
|
|
except Exception:
|
|
pass
|
|
|
|
return result
|
|
|
|
def buddy_ask(self, mission: dict[str, Any] | None, prompt: str) -> dict[str, Any]:
|
|
result = _heuristic_buddy_answer(prompt, mission)
|
|
|
|
if self._mlx_python:
|
|
system = "You are Qwen, the local buddy. Give a concise technical answer."
|
|
user = json.dumps({"mission": mission, "prompt": prompt}, indent=2)
|
|
try:
|
|
result = {
|
|
"answer": self._call_model_text("buddy", system, user),
|
|
"mode": "mlx",
|
|
}
|
|
except Exception:
|
|
pass
|
|
|
|
return result
|