| """Shared runtime context for agents, evaluators, and tools.""" |
| from __future__ import annotations |
|
|
| import argparse |
| import threading |
| from typing import TYPE_CHECKING, Any |
| import numpy as np |
|
|
| if TYPE_CHECKING: |
| from ..evaluator import DefaultEvaluator |
| from ..runtime.workspace import Workspace |
|
|
| _DEFAULT_EVALUATOR = object() |
|
|
|
|
| class AgentContext: |
| """Authoritative runtime context shared across agents, evaluators, and tools. |
| |
| Holds the active dataset arrays (``data``), variable and axis metadata, |
| selected ``target`` and feature names, the active |
| :class:`~sr_harness.evaluator.DefaultEvaluator`, the conversation |
| :class:`~sr_harness.runtime.Workspace`, and cached training/validation views. |
| |
| Args: |
| args: Runtime configuration namespace shared by agents, evaluators, and tools. |
| data: Mapping of variable and axis names to NumPy arrays. |
| target: Name of the active dependent variable, if selected. |
| variable_descriptions: Human-readable descriptions keyed by variable or axis name. |
| variable_axes: Mapping from each variable name to its tuple of axis names. |
| variable_structures: Mapping from edge-aligned variable names to their relation name. |
| relation_names: Set of variable names representing graph or hypergraph edge tables. |
| num_nodes: Total node count when graph or hypergraph relations are present. |
| evaluator: Evaluator instance used for splitting, parameter fitting, and scoring. |
| workspace: Optional managed workspace accessible to workspace-aware tools. |
| """ |
|
|
| def __init__(self, *, args: argparse.Namespace | None = None, data: dict[str, Any] | None = None, target: str | None = None, variable_descriptions: dict[str, str] | None = None, variable_axes: dict[str, tuple[str, ...]] | None = None, variable_structures: dict[str, str] | None = None, relation_names: set[str] | None = None, num_nodes: int | None = None, evaluator: DefaultEvaluator | object = _DEFAULT_EVALUATOR, workspace: Workspace | None = None) -> None: |
| """Initialize an agent runtime context. |
| |
| Args: |
| args: Runtime configuration namespace. |
| data: Mapping of variable and axis names to NumPy arrays. |
| target: Name of the active target variable. |
| variable_descriptions: Descriptions keyed by variable or axis name. |
| variable_axes: Axis tuples keyed by variable name. |
| variable_structures: Relation variable names keyed by edge-aligned variable name. |
| relation_names: Names of graph or hypergraph relation variables. |
| num_nodes: Node count for relational datasets. |
| evaluator: Active formula evaluator instance. |
| workspace: Optional managed workspace instance. |
| """ |
| if args is not None and not isinstance(args, argparse.Namespace): |
| raise TypeError("args must be an argparse.Namespace") |
| self.args = args if args is not None else argparse.Namespace() |
| defaults = { |
| "validation_fraction": 0.0, |
| "split_by": "random", |
| "split_random_state": 42, |
| "split_ood_variable": None, |
| "data_revision": 0, |
| } |
| for name, value in defaults.items(): |
| if not hasattr(self.args, name): |
| setattr(self.args, name, value) |
| self.workspace = workspace |
| if any(not isinstance(name, str) or not name for name in (data or {})): |
| raise TypeError("data keys must be non-empty strings") |
| self.data = {name: np.asarray(value) for name, value in (data or {}).items()} |
| self.target = target |
| if variable_descriptions is None: |
| self.variable_descriptions = {name: "" for name in self.data} |
| else: |
| self.variable_descriptions = dict(variable_descriptions) |
| if variable_axes is None: |
| self.variable_axes = {name: () for name in self.data} |
| else: |
| self.variable_axes = {name: tuple(axes) for name, axes in variable_axes.items()} |
| self.variable_structures = dict(variable_structures or {}) |
| self.relation_names = set(relation_names or self.variable_structures.values()) |
| self.num_nodes = num_nodes |
| from ..evaluator import DefaultEvaluator, GraphEvaluator |
| if evaluator is _DEFAULT_EVALUATOR: |
| evaluator = GraphEvaluator() if self.relation_names else DefaultEvaluator() |
| if not isinstance(evaluator, DefaultEvaluator): |
| raise TypeError("evaluator must be a DefaultEvaluator instance") |
| self.evaluator: DefaultEvaluator = evaluator |
| self._split_cache: dict[str, AgentContext] | None = None |
| self._lock = threading.RLock() |
| self._validate(allow_incomplete=not self.data or self.target is None) |
|
|
| def _validate(self, *, allow_incomplete: bool = False) -> None: |
| names = set(self.data) |
| if not allow_incomplete and self.target not in names: |
| raise ValueError("target must name an entry in data") |
| if set(self.variable_descriptions) != names: |
| raise ValueError("variable_descriptions keys must equal data keys") |
| if any(not isinstance(value, str) for value in self.variable_descriptions.values()): |
| raise TypeError("variable descriptions must be strings") |
| if any( |
| not isinstance(axis, str) or not axis |
| for dimensions in self.variable_axes.values() |
| for axis in dimensions |
| ): |
| raise TypeError("variable_axes values must contain non-empty axis names") |
| non_axes = set(self.variable_axes) |
| axes = {axis for dimensions in self.variable_axes.values() for axis in dimensions} |
| if non_axes & axes or non_axes | axes != names: |
| raise ValueError("variable_axes keys and values must partition data keys") |
| for axis in axes: |
| if self.data[axis].ndim != 1: |
| raise ValueError(f"axis variable {axis!r} must be one-dimensional") |
| if (self.num_nodes is None) != (not self.relation_names): |
| raise ValueError("relation_names and num_nodes must appear together") |
| if self.num_nodes is not None and ( |
| isinstance(self.num_nodes, bool) |
| or not isinstance(self.num_nodes, int) |
| or self.num_nodes < 1 |
| ): |
| raise ValueError("num_nodes must be a positive integer") |
| for relation_name in self.relation_names: |
| if relation_name not in names: |
| raise ValueError("relation_names must reference entries in data") |
| relation = self.data[relation_name] |
| if relation.ndim != 2 or relation.shape[1] not in {2, 3}: |
| raise ValueError(f"relation variable {relation_name!r} must have shape (E, 2) or (H, 3)") |
| if relation.dtype.kind not in "iu": |
| raise ValueError(f"relation variable {relation_name!r} must contain integer endpoints") |
| if np.any(relation < 0) or np.any(relation >= self.num_nodes): |
| raise ValueError(f"relation variable {relation_name!r} endpoints must be in [0, {self.num_nodes})") |
| for variable, structure in self.variable_structures.items(): |
| if variable not in names or structure not in names: |
| raise ValueError("variable_structures must reference entries in data") |
| if variable == structure: |
| raise ValueError("a structured variable cannot reference itself") |
| if structure not in self.relation_names: |
| raise ValueError("variable_structures values must name relation variables") |
| relation = self.data[structure] |
| if relation.ndim != 2 or relation.shape[1] not in {2, 3}: |
| raise ValueError(f"structure variable {structure!r} must have shape (E, 2) or (H, 3)") |
| if self.data[variable].ndim == 0 or self.data[variable].shape[-1] != relation.shape[0]: |
| raise ValueError(f"structured variable {variable!r} must end in the relation dimension") |
| if relation.dtype.kind not in "iu": |
| raise ValueError(f"structure variable {structure!r} must contain integer endpoints") |
| if np.any(relation < 0) or np.any(relation >= self.num_nodes): |
| raise ValueError(f"structure variable {structure!r} endpoints must be in [0, {self.num_nodes})") |
|
|
| def invalidate_splits(self) -> None: |
| """Clear cached training and validation context splits so they are recomputed on next access.""" |
| with self._lock: |
| self._split_cache = None |
|
|
| def _sync_builtin_evaluator(self) -> None: |
| from ..evaluator import DefaultEvaluator, GraphEvaluator |
| if type(self.evaluator) not in {DefaultEvaluator, GraphEvaluator}: |
| return |
| self.evaluator = GraphEvaluator() if self.relation_names else DefaultEvaluator() |
|
|
| def _splits(self) -> dict[str, AgentContext]: |
| with self._lock: |
| if self.target is None: |
| raise ValueError("target must be configured before splitting AgentContext") |
| if self._split_cache is None: |
| splits = self.evaluator.split(self) |
| if set(splits) != {"train", "validation"}: |
| raise ValueError("DefaultEvaluator.split() must return train and validation contexts") |
| if any(not isinstance(split, AgentContext) for split in splits.values()): |
| raise TypeError("DefaultEvaluator.split() values must be AgentContext instances") |
| self._split_cache = splits |
| return self._split_cache |
|
|
| def with_data(self, data: dict[str, np.ndarray]) -> AgentContext: |
| """Create a lightweight view over ``data`` that shares runtime arguments and metadata with this context. |
| |
| Args: |
| data: Replacement mapping of variable and axis arrays. |
| |
| Returns: |
| AgentContext: New context view bound to ``data``. |
| """ |
| return AgentContext( |
| args=self.args, data=data, target=self.target, |
| variable_descriptions={name: self.variable_descriptions[name] for name in data}, |
| variable_axes={name: axes for name, axes in self.variable_axes.items() if name in data}, |
| variable_structures={name: structure for name, structure in self.variable_structures.items() if name in data and structure in data}, |
| relation_names=self.relation_names & set(data), |
| num_nodes=self.num_nodes, evaluator=self.evaluator, workspace=self.workspace, |
| ) |
|
|
| @property |
| def train_split(self) -> AgentContext: |
| """Return the cached training split context, calling ``evaluator.split(self)`` on first access. |
| |
| Returns: |
| AgentContext: Context view containing the training split arrays. |
| """ |
| return self._splits()["train"] |
|
|
| @property |
| def validation_split(self) -> AgentContext: |
| """Return the cached validation split context, calling ``evaluator.split(self)`` on first access. |
| |
| Returns: |
| AgentContext: Context view containing the validation split arrays. |
| """ |
| return self._splits()["validation"] |
|
|
| def axis_names(self) -> tuple[str, ...]: |
| """Return all coordinate axis names referenced by variables in first-appearance order. |
| |
| Returns: |
| tuple[str, ...]: Ordered axis names. |
| """ |
| return tuple(dict.fromkeys(axis for axes in self.variable_axes.values() for axis in axes)) |
|
|
| def variable_names(self) -> tuple[str, ...]: |
| """Return all non-axis variable names in manifest order. |
| |
| Returns: |
| tuple[str, ...]: Ordered variable names. |
| """ |
| return tuple(self.variable_axes) |
|
|
| def feature_names(self) -> tuple[str, ...]: |
| """Return selected input feature names (excluding axes and the active target). |
| |
| Returns: |
| tuple[str, ...]: Active feature variable names. |
| """ |
| return tuple(name for name in self.variable_axes if name != self.target) |
|
|
| def commit_context_data(self, loaded: dict[str, Any]) -> dict[str, Any]: |
| """Replace the live dataset with a validated ``load_context_data()`` bundle and increment ``data_revision``. |
| |
| Args: |
| loaded: Validated dataset dictionary returned by :func:`load_context_data`. |
| |
| Returns: |
| dict[str, Any]: Summary containing the new ``revision``, ``variables``, |
| ``axes``, ``relation_names``, and ``num_nodes``. |
| """ |
| values = {str(name): np.asarray(value) for name, value in loaded["data"].items()} |
| self.data = values |
| self.variable_descriptions = dict(loaded["variable_descriptions"]) |
| self.variable_axes = { |
| name: tuple(axes) for name, axes in loaded["variable_axes"].items() |
| } |
| self.variable_structures = dict(loaded["variable_structures"]) |
| self.relation_names = set(loaded["relation_names"]) |
| self.num_nodes = loaded["num_nodes"] |
| self._sync_builtin_evaluator() |
| self.target = self.target if self.target in values else None |
| self.invalidate_splits() |
| self.args.data_revision += 1 |
| self._validate(allow_incomplete=True) |
| return {"revision": self.args.data_revision, "variables": list(values)} |
|
|
| def commit_data(self, data: dict[str, Any], *, target: str, features: list[str] | None = None, variable_descriptions: dict[str, str] | None = None) -> dict[str, Any]: |
| """Commit aligned 1D/tabular arrays directly into the live context and increment ``data_revision``. |
| |
| Args: |
| data: Mapping of column names to NumPy-compatible arrays. |
| target: Name of the target variable to predict. |
| features: Names of input feature variables to retain. |
| variable_descriptions: Optional human-readable descriptions by column name. |
| |
| Returns: |
| dict[str, Any]: Summary containing the new ``revision``, ``target``, |
| ``features``, ``columns``, and ``rows``. |
| """ |
| selected = list(dict.fromkeys([*(features or []), target])) |
| if target not in data: |
| raise ValueError(f"target column does not exist: {target}") |
| arrays = {name: np.asarray(data[name]) for name in selected} |
| with self._lock: |
| self.data = arrays |
| self.target = target |
| self.variable_descriptions = {name: str((variable_descriptions or {}).get(name, "")) for name in arrays} |
| self.variable_axes = {name: () for name in arrays} |
| self.variable_structures = {} |
| self.relation_names = set() |
| self.num_nodes = None |
| self._sync_builtin_evaluator() |
| self.invalidate_splits() |
| self.args.data_revision += 1 |
| self._validate() |
| return {"revision": self.args.data_revision, "target": target, "features": list(self.feature_names()), "columns": list(arrays), "rows": len(arrays[target])} |
|
|
| def add_features(self, features: dict[str, Any], *, descriptions: dict[str, str] | None = None) -> dict[str, Any]: |
| """Append new sample-aligned feature arrays to the active dataset while keeping the current target. |
| |
| Args: |
| features: Mapping of new feature names to aligned arrays. |
| descriptions: Optional descriptions for the newly added features. |
| |
| Returns: |
| dict[str, Any]: Updated commit summary after merging the new features. |
| |
| Raises: |
| ValueError: If no target variable is currently configured or feature shapes do not align. |
| """ |
| if self.target is None: |
| raise ValueError("target must be configured before adding features") |
| descriptions = self.variable_descriptions | { |
| name: (descriptions or {}).get(name, "") for name in features |
| } |
| return self.commit_data(self.data | features, target=self.target, features=[*self.feature_names(), *features], variable_descriptions=descriptions) |
|
|
| def update_selection(self, *, target: str, features: list[str], variable_descriptions: dict[str, str] | None = None) -> dict[str, Any]: |
| """Update the active target, selected input features, and variable descriptions. |
| |
| Args: |
| target: Name of the target variable. |
| features: List of input feature variable names. |
| variable_descriptions: Optional updated descriptions keyed by variable or axis name. |
| |
| Returns: |
| dict[str, Any]: Summary with the incremented ``revision``, ``target``, and ``features``. |
| """ |
| keep_variables = set(features) | {target} |
| keep_axes = {axis for name, axes in self.variable_axes.items() if name in keep_variables for axis in axes} |
| keep = keep_variables | keep_axes |
| self.data = {name: value for name, value in self.data.items() if name in keep} |
| self.target = target |
| descriptions = self.variable_descriptions | dict(variable_descriptions or {}) |
| self.variable_descriptions = {name: descriptions.get(name, "") for name in self.data} |
| self.variable_axes = {name: axes for name, axes in self.variable_axes.items() if name in keep_variables} |
| self.variable_structures = {name: structure for name, structure in self.variable_structures.items() if name in self.data and structure in self.data} |
| self.relation_names &= set(self.data) |
| if not self.relation_names: |
| self.num_nodes = None |
| self._sync_builtin_evaluator() |
| self.invalidate_splits() |
| self.args.data_revision += 1 |
| self._validate() |
| return {"revision": self.args.data_revision, "target": target, "features": list(self.feature_names()), "selection_changed": True} |
|
|
| def schema(self) -> dict[str, Any]: |
| """Return a JSON-serializable summary of variables, shapes, dtypes, axes, and roles in the context. |
| |
| Returns: |
| dict[str, Any]: Schema dictionary suitable for WebUI previews and diagnostics. |
| """ |
| axes = set(self.axis_names()) |
| return { |
| "revision": self.args.data_revision, "num_nodes": self.num_nodes, |
| "target": self.target, "features": list(self.feature_names()), "columns": list(self.data), |
| "rows": len(self.data[self.target]) if self.target in self.data else 0, |
| "variable_descriptions": dict(self.variable_descriptions), |
| "variables": {name: {"shape": list(value.shape), "dtype": str(value.dtype), "axes": list(self.variable_axes.get(name, ())), "description": self.variable_descriptions[name], **({"kind": "relation"} if name in self.relation_names else {})} for name, value in self.data.items() if name not in axes}, |
| "axes": {name: {"size": len(self.data[name]), "dtype": str(self.data[name].dtype), "description": self.variable_descriptions[name]} for name in axes}, |
| } |
|
|
| def __getstate__(self) -> dict[str, Any]: |
| state = self.__dict__.copy() |
| state.pop("_lock", None) |
| return state |
|
|
| def __setstate__(self, state: dict[str, Any]) -> None: |
| self.__dict__.update(state) |
| self._lock = threading.RLock() |
|
|
|
|
