local-rig/local_rig/rig_schema.py

"""rig.yaml 解析 + 校验（借鉴 OpenRig pods/members/edges schema）。

schema:
  version: "0.1"
  name: my-rig
  pods:
    - id: dev
      label: Development
      members:
        - id: lead
          runtime: hermes-cli | terminal | hermes-acp
          target: hermes          # hermes-* 运行时用
          command: "python3 qa.py"  # terminal 运行时用
      edges:
        - kind: delegates_to | can_observe | collaborates_with
          from: lead
          to: qa
"""
from __future__ import annotations

import re
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import yaml

VALID_RUNTIMES = {"terminal", "hermes-cli", "hermes-acp"}
VALID_EDGE_KINDS = {"delegates_to", "can_observe", "collaborates_with"}
_ID_RE = re.compile(r"^[a-z][a-z0-9_-]{0,31}$")


class RigSchemaError(ValueError):
    """rig.yaml 格式/语义错误。"""


@dataclass
class Member:
    id: str
    runtime: str
    target: str | None = None       # hermes-* 运行时
    command: str | None = None      # terminal 运行时
    profile: str = "default"
    cwd: str = "."
    # 【dsh@mac13 补】真机实测：hermes oneshot 不带 `-m` 时服务端直接报
    #   `HTTP 400: The supported API model names are deepseek-flash, deepseek-v4-pro, but you passed .`
    # ⇒ 模型必须显式传。默认值取本机 provider 支持的名字，可在 rig.yaml 覆盖。
    model: str = "deepseek-flash"

    def __post_init__(self):
        if self.runtime not in VALID_RUNTIMES:
            raise RigSchemaError(f"member '{self.id}': runtime '{self.runtime}' 不合法，可选 {sorted(VALID_RUNTIMES)}")
        if self.runtime == "terminal" and not self.command:
            raise RigSchemaError(f"member '{self.id}': terminal 运行时必须指定 command")
        if self.runtime in ("hermes-cli", "hermes-acp") and not self.target:
            raise RigSchemaError(f"member '{self.id}': {self.runtime} 运行时必须指定 target")


@dataclass
class Edge:
    kind: str
    from_id: str
    to_id: str

    def __post_init__(self):
        if self.kind not in VALID_EDGE_KINDS:
            raise RigSchemaError(f"edge kind '{self.kind}' 不合法，可选 {sorted(VALID_EDGE_KINDS)}")


@dataclass
class Pod:
    id: str
    label: str = ""
    members: list[Member] = field(default_factory=list)
    edges: list[Edge] = field(default_factory=list)


@dataclass
class Rig:
    name: str
    pods: list[Pod] = field(default_factory=list)
    version: str = "0.1"

    def all_members(self) -> dict[str, Member]:
        """扁平化所有 member，key = 'pod_id.member_id'。"""
        out: dict[str, Member] = {}
        for pod in self.pods:
            for m in pod.members:
                out[f"{pod.id}.{m.id}"] = m
        return out

    def resolve_edge(self, e: Edge) -> tuple[str, str]:
        """把 from/to 解析为 'pod.member' 全名。"""
        return self._resolve(e.from_id), self._resolve(e.to_id)

    def _resolve(self, ref: str) -> str:
        if "." in ref:
            pod_id, _, member_id = ref.partition(".")
            for pod in self.pods:
                if pod.id == pod_id:
                    for m in pod.members:
                        if m.id == member_id:
                            return f"{pod_id}.{member_id}"
            raise RigSchemaError(f"edge 引用 '{ref}' 找不到对应 member")
        # 无 pod 前缀：全局查找
        for pod in self.pods:
            for m in pod.members:
                if m.id == ref:
                    return f"{pod.id}.{m.id}"
        raise RigSchemaError(f"edge 引用 '{ref}' 找不到对应 member")


def _check_id(id_val: str, ctx: str):
    if not _ID_RE.match(id_val or ""):
        raise RigSchemaError(f"{ctx} id '{id_val}' 不合法（须小写字母开头，≤32 字符）")


def parse_rig(path: str | Path) -> Rig:
    """解析 rig.yaml 并校验。"""
    p = Path(path)
    if not p.exists():
        raise RigSchemaError(f"rig 文件不存在: {p}")
    data = yaml.safe_load(p.read_text(encoding="utf-8"))
    if not isinstance(data, dict):
        raise RigSchemaError("rig.yaml 顶层必须是 mapping")

    name = data.get("name")
    if not name or not isinstance(name, str):
        raise RigSchemaError("rig.yaml 缺少 name（字符串）")

    pods_raw = data.get("pods") or []
    if not isinstance(pods_raw, list) or not pods_raw:
        raise RigSchemaError("rig.yaml 缺少 pods（非空列表）")

    pods: list[Pod] = []
    pod_ids: set[str] = set()

    for i, pr in enumerate(pods_raw):
        if not isinstance(pr, dict):
            raise RigSchemaError(f"pods[{i}] 必须是 mapping")
        pid = pr.get("id")
        _check_id(pid, f"pods[{i}]")
        if pid in pod_ids:
            raise RigSchemaError(f"重复 pod id: {pid}")
        pod_ids.add(pid)

        members: list[Member] = []
        member_ids: set[str] = set()
        for j, mr in enumerate(pr.get("members") or []):
            if not isinstance(mr, dict):
                raise RigSchemaError(f"pod '{pid}' members[{j}] 必须是 mapping")
            mid = mr.get("id")
            _check_id(mid, f"pod '{pid}' members[{j}]")
            if mid in member_ids:
                raise RigSchemaError(f"pod '{pid}' 重复 member id: {mid}")
            member_ids.add(mid)
            members.append(Member(
                id=mid,
                runtime=mr.get("runtime", "terminal"),
                target=mr.get("target"),
                command=mr.get("command"),
                profile=mr.get("profile", "default"),
                cwd=mr.get("cwd", "."),
                model=mr.get("model", "deepseek-flash"),   # 【dsh@mac13 补】可在 rig.yaml 覆盖
            ))

        edges: list[Edge] = []
        for k, er in enumerate(pr.get("edges") or []):
            if not isinstance(er, dict):
                raise RigSchemaError(f"pod '{pid}' edges[{k}] 必须是 mapping")
            edges.append(Edge(
                kind=er.get("kind", "delegates_to"),
                from_id=er.get("from", ""),
                to_id=er.get("to", ""),
            ))

        pods.append(Pod(id=pid, label=pr.get("label", ""), members=members, edges=edges))

    rig = Rig(name=name, pods=pods, version=str(data.get("version", "0.1")))

    # 全局校验：edge 引用必须存在
    all_ids = set(rig.all_members().keys())
    for pod in pods:
        for e in pod.edges:
            try:
                rig.resolve_edge(e)
            except RigSchemaError:
                raise
            # 确保 from/to 在已知 member 里
            for ref in (e.from_id, e.to_id):
                resolved = rig._resolve(ref)
                if resolved not in all_ids:
                    raise RigSchemaError(f"pod '{pod.id}' edge 引用 '{ref}' → '{resolved}' 不是已知 member")

    return rig


def load_rig(path: str | Path) -> Rig:
    """parse_rig 的别名（语义更清晰）。"""
    return parse_rig(path)