open-standards/jouleclaw/jouleclaw-rs/crates/jouleclaw-hydra-client/clients/python/jclaw_hydra/client.py

"""HydraClient — parse + validate + dispatch + emit CallReport."""

from __future__ import annotations

import json
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional, Protocol

try:
    from blake3 import blake3 as _blake3
except ImportError as _e:  # pragma: no cover - import-time only
    raise ImportError(
        "the `blake3` package is required: pip install blake3"
    ) from _e

from .hydra import ApiDocumentation, Operation, parse_api
from .shacl import ShapeGraph, ValidationReport, validate as shacl_validate


class ClientError(Exception):
    """Base class for HydraClient errors."""


class OperationNotFound(ClientError):
    pass


class InputInvalid(ClientError):
    def __init__(self, shape: str, report: ValidationReport) -> None:
        super().__init__(
            f"input does not validate against shape `{shape}` "
            f"({len(report.violations)} violation(s))"
        )
        self.shape = shape
        self.report = report


GRADES = ("HwShunt", "ModelBased", "Estimator")


@dataclass
class TransportResponse:
    status: int
    body: Any
    # The transport's figure for the dispatch's energy, if it has one.
    # None: none, and the client emits a zero graded "Estimator".
    joules_uj: Optional[int] = None
    # The grade of the meter that produced `joules_uj`: "HwShunt" only for a
    # hardware reading, "ModelBased", or "Estimator". None (or any other
    # value) is sealed "Estimator": the client cannot tell how a transport
    # obtained a figure, and a figure must not claim better provenance than
    # how it was obtained.
    joules_provenance: Optional[str] = None

    def graded_joules(self) -> "tuple[int, str]":
        """The figure and grade the client seals: the transport's figure at
        its stated grade, an ungraded figure as an estimate, and no figure as
        a zero estimate. Matches `TransportResponse::graded_joules` (Rust)."""
        if self.joules_uj is None:
            return 0, "Estimator"
        if self.joules_provenance in GRADES:
            return self.joules_uj, self.joules_provenance
        return self.joules_uj, "Estimator"


class Transport(Protocol):
    def dispatch(self, method: str, operation_id: str, input_value: Any) -> TransportResponse: ...


class MockTransport:
    """Reference transport — echo-mode by default; canned responses
    by operation id otherwise. Mirrors `jouleclaw_hydra_client::MockTransport`."""

    def __init__(self) -> None:
        self._responses: Dict[str, TransportResponse] = {}

    def respond(self, op_id: str, response: TransportResponse) -> "MockTransport":
        self._responses[op_id] = response
        return self

    def dispatch(self, _method: str, operation_id: str, input_value: Any) -> TransportResponse:
        if operation_id in self._responses:
            return self._responses[operation_id]
        # Echo-mode — matches jouleclaw_hydra_client::MockTransport byte
        # for byte (only "echo" in the body; no "op" sidecar).
        return TransportResponse(
            status=200,
            body={"echo": input_value},
            joules_uj=None,
        )


@dataclass
class CallReport:
    tool_id: str
    status: int
    input_hash: str
    output_hash: str
    status_matches_contract: bool
    body: Any
    touch: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "tool_id": self.tool_id,
            "status": self.status,
            "input_hash": self.input_hash,
            "output_hash": self.output_hash,
            "status_matches_contract": self.status_matches_contract,
            "body": self.body,
            "touch": self.touch,
        }


class HydraClient:
    """Generic third-generation agent client. Same surface as the Rust
    `HydraClient` and the Go reference client; produces byte-equal
    BLAKE3 input/output hashes on the same vectors."""

    def __init__(self, api: ApiDocumentation, shapes: Optional[ShapeGraph] = None) -> None:
        self.api = api
        self.shapes = shapes or ShapeGraph()

    @classmethod
    def from_json_ld(cls, source: Any, shapes: Optional[ShapeGraph] = None) -> "HydraClient":
        return cls(parse_api(source), shapes)

    def operations(self) -> List[Operation]:
        result: List[Operation] = []
        for c in self.api.supported_class:
            result.extend(c.supported_operation)
        return result

    def find_operation(self, name: str) -> Optional[Operation]:
        for op in self.operations():
            if op.jc_tool_id == name or op.id == name or op.title == name:
                return op
        return None

    def validate_input(self, op: Operation, value: Any) -> None:
        if op.expects.endswith("#Nothing"):
            return
        shape = self.shapes.get(op.expects)
        if shape is None:
            return  # open-world default
        report = shacl_validate(self.shapes, op.expects, value)
        if not report.conforms:
            raise InputInvalid(op.expects, report)

    def call(self, op_name: str, input_value: Any, transport: Transport) -> CallReport:
        op = self.find_operation(op_name)
        if op is None:
            raise OperationNotFound(f"operation `{op_name}` not found")
        self.validate_input(op, input_value)
        response = transport.dispatch(op.method, op.id, input_value)
        tool_id = op.jc_tool_id or op.id
        input_hash = blake3_hex(canonical_json(input_value))
        output_hash = blake3_hex(canonical_json(response.body))
        status_matches = any(
            s.get("code") == response.status for s in op.status_codes
        )
        joules_uj, provenance = response.graded_joules()
        return CallReport(
            tool_id=tool_id,
            status=response.status,
            input_hash=input_hash,
            output_hash=output_hash,
            status_matches_contract=status_matches,
            body=response.body,
            touch={
                "tool_id": tool_id,
                "joules_uj": joules_uj,
                "energy_provenance": provenance,
            },
        )


# ─────────────────────────────────────────────────────────────────────
# Canonical-JSON + BLAKE3 — the bytes the three implementations agree on
# ─────────────────────────────────────────────────────────────────────


def canonical_json(value: Any) -> str:
    """Lexically-sorted, no-whitespace JSON. Matches the Rust client's
    canonical form so the BLAKE3 hash is identical across implementations.
    """
    return json.dumps(_sort_keys(value), separators=(",", ":"), ensure_ascii=False)


def blake3_hex(text: str) -> str:
    return _blake3(text.encode("utf-8")).hexdigest()


def _sort_keys(value: Any) -> Any:
    if isinstance(value, dict):
        return {k: _sort_keys(value[k]) for k in sorted(value.keys())}
    if isinstance(value, list):
        return [_sort_keys(v) for v in value]
    return value