| """Manifest-backed NumPy data stored in a workspace directory.""" |
|
|
| from __future__ import annotations |
|
|
| import json |
| import os |
| from pathlib import Path |
| import tempfile |
| from typing import TYPE_CHECKING, Any |
|
|
| import numpy as np |
|
|
| if TYPE_CHECKING: |
| from ..runtime.workspace import Workspace |
|
|
|
|
| class ContextManifestError(ValueError): |
| """Raised when a context-data manifest cannot be validated.""" |
|
|
|
|
| class _ContextDataReader: |
| """Private implementation shared by the public loading functions.""" |
|
|
| MANIFEST_NAME = "manifest.json" |
| REQUIRED_ROOT_FIELDS = {"variables", "axes"} |
| OPTIONAL_ROOT_FIELDS = {"num_nodes"} |
| REQUIRED_VARIABLE_FIELDS = {"file", "description", "axes"} |
| OPTIONAL_VARIABLE_FIELDS = {"kind", "structure"} |
| AXIS_SOURCE_FIELDS = {"values", "file", "size"} |
| AXIS_FIELDS = {"description", *AXIS_SOURCE_FIELDS} |
|
|
| def __init__(self, directory: str | Path, *, workspace: Workspace | None = None): |
| self.directory = Path(directory).resolve() |
| self.workspace = workspace |
|
|
| def inspect(self) -> dict[str, Any]: |
| """Validate the store and return diagnostics without raising. |
| |
| Returns: |
| A serializable report containing errors, warnings, and array summaries. |
| """ |
| errors: list[str] = [] |
| warnings: list[str] = [] |
| manifest = self._read_manifest(errors) |
| loaded = ( |
| self._validate_manifest(manifest, errors, warnings) |
| if manifest is not None |
| else None |
| ) |
| report: dict[str, Any] = { |
| "valid": not errors, |
| "path": str(self.directory), |
| "errors": errors, |
| "warnings": warnings, |
| "num_nodes": loaded["num_nodes"] if loaded is not None else None, |
| "variables": {}, |
| "axes": {}, |
| } |
| if loaded is not None: |
| report["variables"] = { |
| name: { |
| "file": manifest["variables"][name]["file"], |
| "shape": list(value.shape), |
| "dtype": str(value.dtype), |
| "axes": list(loaded["variable_axes"][name]), |
| **( |
| {"structure": loaded["variable_structures"][name]} |
| if name in loaded["variable_structures"] else {} |
| ), |
| **({"kind": "relation"} if name in loaded["relation_names"] else {}), |
| "description": loaded["variable_descriptions"][name], |
| } |
| for name, value in loaded["data"].items() |
| if name in manifest["variables"] |
| } |
| report["axes"] = { |
| name: { |
| "size": len(loaded["data"][name]), |
| "dtype": str(loaded["data"][name].dtype), |
| "storage": axis["storage"], |
| "description": loaded["variable_descriptions"][name], |
| } |
| for name, axis in loaded["axis_metadata"].items() |
| } |
| return report |
|
|
| def load(self) -> dict[str, Any]: |
| """Validate and load all variables and axes. |
| |
| Returns: |
| An AgentContext-ready mapping containing arrays and metadata. |
| |
| Raises: |
| ContextManifestError: If the manifest or referenced arrays are invalid. |
| """ |
| errors: list[str] = [] |
| warnings: list[str] = [] |
| manifest = self._read_manifest(errors) |
| loaded = ( |
| self._validate_manifest(manifest, errors, warnings) |
| if manifest is not None else None |
| ) |
| if errors or loaded is None: |
| raise ContextManifestError("Invalid context.data:\n- " + "\n- ".join(errors)) |
| return loaded |
|
|
| def _read_manifest(self, errors: list[str]) -> dict[str, Any] | None: |
| manifest_path = self.directory / self.MANIFEST_NAME |
| if not self.directory.is_dir(): |
| errors.append(f"directory does not exist: {self.directory}") |
| return None |
| if not manifest_path.is_file(): |
| errors.append(f"missing {self.MANIFEST_NAME}") |
| return None |
|
|
| def reject_duplicates(pairs): |
| result = {} |
| for key, value in pairs: |
| if key in result: |
| raise ContextManifestError(f"duplicate JSON key: {key}") |
| result[key] = value |
| return result |
|
|
| try: |
| opener = self.workspace.safe_open if self.workspace is not None else open |
| with opener(manifest_path, "r", encoding="utf-8") as f: |
| manifest = json.load(f, object_pairs_hook=reject_duplicates) |
| except (ContextManifestError, json.JSONDecodeError, OSError) as exc: |
| errors.append(f"cannot read {self.MANIFEST_NAME}: {exc}") |
| return None |
| if not isinstance(manifest, dict): |
| errors.append("manifest root must be a JSON object") |
| return None |
| return manifest |
|
|
| def _validate_manifest( |
| self, |
| manifest: dict[str, Any], |
| errors: list[str], |
| warnings: list[str], |
| ) -> dict[str, Any] | None: |
| self._check_fields( |
| "manifest", |
| manifest, |
| self.REQUIRED_ROOT_FIELDS, |
| errors, |
| optional=self.OPTIONAL_ROOT_FIELDS, |
| ) |
| variables = manifest.get("variables") |
| axes = manifest.get("axes") |
| num_nodes = manifest.get("num_nodes") |
| if "num_nodes" in manifest and ( |
| isinstance(num_nodes, bool) or not isinstance(num_nodes, int) or num_nodes < 1 |
| ): |
| errors.append("manifest.num_nodes must be a positive integer") |
| num_nodes = None |
| if not isinstance(variables, dict) or not variables: |
| errors.append("manifest.variables must be a non-empty object") |
| variables = {} |
| if not isinstance(axes, dict): |
| errors.append("manifest.axes must be an object") |
| axes = {} |
|
|
| loaded_axes: dict[str, dict[str, Any]] = {} |
| referenced_files = {self.MANIFEST_NAME} |
| for name, spec in axes.items(): |
| location = f"axes.{name}" |
| if not self._valid_name(name): |
| errors.append(f"{location}: name must be non-empty and contain no path separators") |
| continue |
| if not isinstance(spec, dict): |
| errors.append(f"{location} must be an object") |
| continue |
| if "description" not in spec: |
| errors.append(f"{location} is missing fields: ['description']") |
| if extra := sorted(set(spec) - self.AXIS_FIELDS): |
| errors.append(f"{location} contains unsupported fields: {extra}") |
| description = self._description(location, spec, errors) |
| sources = self.AXIS_SOURCE_FIELDS & set(spec) |
| if len(sources) != 1: |
| errors.append(f"{location} must contain exactly one of values, file, or size") |
| continue |
| source = next(iter(sources)) |
| try: |
| if source == "values": |
| if not isinstance(spec["values"], list) or not spec["values"]: |
| raise ValueError("values must be a non-empty JSON array") |
| if any(isinstance(value, (list, dict)) for value in spec["values"]): |
| raise ValueError("values must contain scalar JSON values") |
| values = np.asarray(spec["values"]) |
| elif source == "size": |
| size = spec["size"] |
| if isinstance(size, bool) or not isinstance(size, int) or size < 1: |
| raise ValueError("size must be a positive integer") |
| values = np.arange(size) |
| else: |
| expected = f"{name}.npy" |
| if spec["file"] != expected: |
| raise ValueError(f"file must be {expected!r}") |
| referenced_files.add(expected) |
| values = self._load_array(expected) |
| if values.ndim != 1: |
| raise ValueError(f"axis array must be one-dimensional, got shape {values.shape}") |
| if values.ndim != 1: |
| raise ValueError("values must be one-dimensional") |
| loaded_axes[name] = {"values": values, "description": description, "storage": source} |
| except (OSError, ValueError) as exc: |
| errors.append(f"{location}: {exc}") |
|
|
| loaded_values: dict[str, np.ndarray] = {} |
| variable_axes: dict[str, tuple[str, ...]] = {} |
| variable_structures: dict[str, str] = {} |
| relation_names: set[str] = set() |
| descriptions: dict[str, str] = {} |
| used_axes: set[str] = set() |
| for name, spec in variables.items(): |
| location = f"variables.{name}" |
| if not self._valid_name(name): |
| errors.append(f"{location}: name must be non-empty and contain no path separators") |
| continue |
| if not isinstance(spec, dict): |
| errors.append(f"{location} must be an object") |
| continue |
| self._check_fields( |
| location, |
| spec, |
| self.REQUIRED_VARIABLE_FIELDS, |
| errors, |
| optional=self.OPTIONAL_VARIABLE_FIELDS, |
| ) |
| description = self._description(location, spec, errors) |
| structure = spec.get("structure") |
| if structure is not None and not self._valid_name(structure): |
| errors.append(f"{location}.structure must name an A or T variable") |
| structure = None |
| kind = spec.get("kind") |
| if kind not in {None, "relation"}: |
| errors.append(f"{location}.kind must be 'relation' when provided") |
| kind = None |
| expected = f"{name}.npy" |
| if spec.get("file") != expected: |
| errors.append(f"{location}.file must be {expected!r}") |
| continue |
| declared_axes = spec.get("axes") |
| if not isinstance(declared_axes, list) or any( |
| not isinstance(axis, str) or not axis for axis in declared_axes |
| ): |
| errors.append(f"{location}.axes must be an array of non-empty axis names") |
| continue |
| missing_axes = [axis for axis in declared_axes if axis not in axes] |
| if missing_axes: |
| errors.append(f"{location}.axes references missing axes: {missing_axes}") |
| continue |
| referenced_files.add(expected) |
| try: |
| value = self._load_array(expected) |
| except (OSError, ValueError) as exc: |
| errors.append(f"{location}: {exc}") |
| continue |
| if value.ndim != len(declared_axes): |
| errors.append( |
| f"{location}: array has {value.ndim} dimensions but declares " |
| f"{len(declared_axes)} axes" |
| ) |
| continue |
| mismatches = [ |
| f"{axis}={len(loaded_axes[axis]['values'])} (data={value.shape[index]})" |
| for index, axis in enumerate(declared_axes) |
| if axis in loaded_axes and len(loaded_axes[axis]["values"]) != value.shape[index] |
| ] |
| if mismatches: |
| errors.append(f"{location}: axis length mismatch: {', '.join(mismatches)}") |
| continue |
| loaded_values[name] = value |
| variable_axes[name] = tuple(declared_axes) |
| if structure is not None: |
| variable_structures[name] = structure |
| if kind == "relation": |
| relation_names.add(name) |
| descriptions[name] = description |
| used_axes.update(declared_axes) |
|
|
| unused_axes = sorted(set(axes) - used_axes) |
| if unused_axes: |
| errors.append(f"manifest.axes contains unreferenced axes: {unused_axes}") |
|
|
| if bool(relation_names) != (num_nodes is not None): |
| errors.append("manifest.num_nodes and relation metadata must appear together") |
| for relation_name in relation_names: |
| location = f"variables.{relation_name}" |
| relation = loaded_values.get(relation_name) |
| if relation is None: |
| continue |
| if relation.ndim != 2 or relation.shape[1] not in {2, 3}: |
| errors.append( |
| f"{location}: relation must have shape (E, 2) or (H, 3), got {relation.shape}" |
| ) |
| continue |
| if relation.dtype.kind not in "iu": |
| errors.append(f"{location}: relation endpoints must use an integer dtype") |
| elif num_nodes is not None and (np.any(relation < 0) or np.any(relation >= num_nodes)): |
| errors.append(f"{location}: relation endpoints must be in [0, {num_nodes})") |
| for variable, structure in variable_structures.items(): |
| location = f"variables.{variable}.structure" |
| if variable == structure: |
| errors.append(f"{location}: a structured variable cannot reference itself") |
| continue |
| if structure not in loaded_values: |
| errors.append(f"{location} references missing variable {structure!r}") |
| continue |
| if structure not in relation_names: |
| errors.append(f"{location} must reference a variable with kind 'relation'") |
| continue |
| relation = loaded_values[structure] |
| if relation.ndim != 2 or relation.shape[1] not in {2, 3}: |
| errors.append( |
| f"{location}: {structure!r} must have shape (E, 2) or (H, 3), " |
| f"got {relation.shape}" |
| ) |
| continue |
| dependent = loaded_values.get(variable) |
| if dependent is not None and ( |
| dependent.ndim == 0 or dependent.shape[-1] != relation.shape[0] |
| ): |
| errors.append( |
| f"{location}: {variable!r} must end in a dimension of " |
| f"length {relation.shape[0]}" |
| ) |
| if relation.dtype.kind not in "iu": |
| errors.append(f"{location}: relation endpoints must use an integer dtype") |
| elif num_nodes is not None and (np.any(relation < 0) or np.any(relation >= num_nodes)): |
| errors.append(f"{location}: relation endpoints must be in [0, {num_nodes})") |
| actual_npy = {path.name for path in self.directory.glob("*.npy") if path.is_file()} |
| untracked = sorted(actual_npy - referenced_files) |
| if untracked: |
| warnings.append(f"unreferenced NPY files: {untracked}") |
| if errors: |
| return None |
| data = {**{name: axis["values"] for name, axis in loaded_axes.items()}, **loaded_values} |
| all_descriptions = {name: axis["description"] for name, axis in loaded_axes.items()} | descriptions |
| return { |
| "data": data, |
| "variable_descriptions": all_descriptions, |
| "variable_axes": variable_axes, |
| "variable_structures": variable_structures, |
| "relation_names": relation_names, |
| "num_nodes": num_nodes, |
| "axis_metadata": loaded_axes, |
| } |
|
|
| def _load_array(self, filename: str) -> np.ndarray: |
| path = self.directory / filename |
| if not path.is_file() or path.parent != self.directory: |
| raise ValueError(f"missing array file: {filename}") |
| opener = self.workspace.safe_open if self.workspace is not None else open |
| with opener(path, "rb") as f: |
| value = np.load(f, allow_pickle=False) |
| if not isinstance(value, np.ndarray): |
| raise ValueError(f"{filename} does not contain one NPY array") |
| if value.dtype.hasobject: |
| raise ValueError(f"{filename} uses object dtype, which is not allowed") |
| return value |
|
|
| @staticmethod |
| def _valid_name(name: Any) -> bool: |
| return ( |
| isinstance(name, str) |
| and bool(name.strip()) |
| and name not in {".", ".."} |
| and "/" not in name |
| and "\\" not in name |
| and "\0" not in name |
| ) |
|
|
| @staticmethod |
| def _description(location: str, spec: dict[str, Any], errors: list[str]) -> str: |
| description = spec.get("description") |
| if not isinstance(description, str): |
| errors.append(f"{location}.description must be a string") |
| return "" |
| return description.strip() |
|
|
| @staticmethod |
| def _check_fields( |
| location: str, |
| value: dict[str, Any], |
| expected: set[str], |
| errors: list[str], |
| *, |
| optional: set[str] | None = None, |
| ) -> None: |
| optional = optional or set() |
| missing = sorted(expected - set(value)) |
| extra = sorted(set(value) - expected - optional) |
| if missing: |
| errors.append(f"{location} is missing fields: {missing}") |
| if extra: |
| errors.append(f"{location} contains unsupported fields: {extra}") |
|
|
|
|
| def load_context_data(directory: str | Path, *, workspace: Workspace | None = None) -> dict[str, Any]: |
| """Validate and load one manifest-backed ``context.data`` directory. |
| |
| Args: |
| directory: Directory containing ``manifest.json`` and array files. |
| workspace: Workspace enforcing safe reads, when available. |
| |
| Returns: |
| An AgentContext-ready mapping containing arrays and metadata. |
| """ |
| return _ContextDataReader(directory, workspace=workspace).load() |
|
|
|
|
| def inspect_context_data(directory: str | Path, *, workspace: Workspace | None = None) -> dict[str, Any]: |
| """Return validation diagnostics without raising for manifest errors. |
| |
| Args: |
| directory: Directory containing ``manifest.json`` and array files. |
| workspace: Workspace enforcing safe reads, when available. |
| |
| Returns: |
| A serializable validation report. |
| """ |
| return _ContextDataReader(directory, workspace=workspace).inspect() |
|
|
|
|
| def update_context_data_descriptions( |
| directory: str | Path, |
| descriptions: dict[str, str], |
| *, |
| workspace: Workspace | None = None, |
| ) -> dict[str, Any]: |
| """Update variable and axis descriptions in ``manifest.json``. |
| |
| The complete store is validated before the edit. An atomic replacement is |
| used when the directory permits creating entries. If the directory is |
| read-only but ``manifest.json`` itself is writable, the existing file is |
| updated in place, matching normal POSIX file-permission semantics. Because |
| only existing string-valued description fields are changed, structural |
| validity is preserved. The returned mapping matches :func:`load_context_data`. |
| """ |
| directory = Path(directory).resolve() |
| manifest_path = directory / _ContextDataReader.MANIFEST_NAME |
| loaded = load_context_data(directory, workspace=workspace) |
| opener = workspace.safe_open if workspace is not None else open |
| if not isinstance(descriptions, dict) or any( |
| not isinstance(name, str) or not isinstance(value, str) |
| for name, value in descriptions.items() |
| ): |
| raise ValueError("descriptions must map variable names to text") |
| unknown = set(descriptions) - set(loaded["data"]) |
| if unknown: |
| raise ValueError(f"unknown context.data variables: {sorted(unknown)}") |
| with opener(manifest_path, "r", encoding="utf-8") as f: |
| manifest = json.load(f) |
| for section in ("variables", "axes"): |
| for name, spec in manifest[section].items(): |
| if name in descriptions: |
| spec["description"] = descriptions[name].strip() |
| serialized = json.dumps(manifest, ensure_ascii=False, indent=2) + "\n" |
| if workspace is not None: |
| relative_manifest = str(manifest_path.relative_to(workspace.path)) |
| temporary_path = None |
| try: |
| with tempfile.NamedTemporaryFile( |
| mode="w", |
| encoding="utf-8", |
| dir=workspace.path, |
| prefix=".manifest-", |
| suffix=".json.tmp", |
| delete=False, |
| ) as stream: |
| temporary_path = Path(stream.name) |
| stream.write(serialized) |
| stream.flush() |
| os.fsync(stream.fileno()) |
| workspace.replace_file(temporary_path, relative_manifest) |
| temporary_path = None |
| finally: |
| if temporary_path is not None: |
| temporary_path.unlink(missing_ok=True) |
| updated = dict(loaded) |
| updated["variable_descriptions"] = loaded["variable_descriptions"] | { |
| name: value.strip() for name, value in descriptions.items() |
| } |
| return updated |
| temporary_path = None |
| try: |
| try: |
| with tempfile.NamedTemporaryFile( |
| mode="w", |
| encoding="utf-8", |
| dir=directory, |
| prefix=".manifest-", |
| suffix=".json.tmp", |
| delete=False, |
| ) as stream: |
| temporary_path = Path(stream.name) |
| stream.write(serialized) |
| stream.flush() |
| os.fsync(stream.fileno()) |
| os.replace(temporary_path, manifest_path) |
| temporary_path = None |
| except PermissionError: |
| if temporary_path is not None: |
| temporary_path.unlink(missing_ok=True) |
| with manifest_path.open("w", encoding="utf-8") as stream: |
| stream.write(serialized) |
| stream.flush() |
| os.fsync(stream.fileno()) |
| finally: |
| if temporary_path is not None: |
| temporary_path.unlink(missing_ok=True) |
| updated = dict(loaded) |
| updated["variable_descriptions"] = loaded["variable_descriptions"] | { |
| name: value.strip() for name, value in descriptions.items() |
| } |
| return updated |
|
|
