Upload analysis/trajectory_extractor.py

analysis/trajectory_extractor.py

@@ -0,0 +1,915 @@
+"""
+analysis.trajectory_extractor
+=============================
+
+Convert raw WorldSmithAI simulation history into structured trajectory data.
+
+This module is intentionally placed above the simulation engine. It consumes
+world history, metric logs, and event logs without mutating runtime agents,
+resources, behaviors, policies, schedulers, or worlds.
+
+The extractor is designed to be tolerant of several history formats because the
+core engine may evolve independently. It accepts either:
+
+    - a world-like object with ``history``
+    - a raw sequence of snapshot mappings
+    - a single snapshot mapping
+
+The output is a ``TrajectoryData`` dataclass consumed by later analysis modules.
+
+Example
+-------
+
+    extractor = TrajectoryExtractor()
+    trajectory = extractor.extract(world)
+
+    print(trajectory.step_count)
+    print(trajectory.agent_type_counts)
+    print(trajectory.population_timeseries)
+
+Architecture
+------------
+
+    world.history
+        ↓
+    TrajectoryExtractor.extract(...)
+        ↓
+    TrajectoryData
+        ↓
+    simulation_analyzer / report_generator / scenario comparator
+"""
+
+from __future__ import annotations
+
+import copy
+import logging
+import math
+from collections import defaultdict
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass, field
+from typing import Any, Protocol
+
+
+logger = logging.getLogger(__name__)
+
+
+class WorldHistoryProtocol(Protocol):
+    """Structural protocol for world-like objects with history."""
+
+    history: Any
+    step_count: int
+
+
+@dataclass(frozen=True)
+class TrajectoryExtractionConfig:
+    """Configuration for trajectory extraction.
+
+    Attributes:
+        count_alive_only:
+            If true, population counts exclude agents with ``alive=False``.
+        include_snapshot_events:
+            If true, event-like records inside each snapshot are added to the
+            normalized event log.
+        include_snapshot_metrics:
+            If true, metric-like records inside each snapshot are added to the
+            normalized metric history.
+        preserve_unknown_snapshot_fields:
+            If true, selected unrecognized top-level snapshot keys are preserved
+            in metadata for debugging.
+        max_preserved_unknown_values:
+            Maximum number of unknown snapshot values preserved in metadata.
+    """
+
+    count_alive_only: bool = True
+    include_snapshot_events: bool = True
+    include_snapshot_metrics: bool = True
+    preserve_unknown_snapshot_fields: bool = False
+    max_preserved_unknown_values: int = 20
+
+
+@dataclass(frozen=True)
+class TrajectoryData:
+    """Structured simulation trajectory data.
+
+    Attributes:
+        population_timeseries:
+            Mapping from step to counts by agent type.
+        resource_timeseries:
+            Mapping from step to total resource amount by resource type.
+        event_log:
+            Normalized event records. Each record should include a ``step`` key
+            when the step is known.
+        metric_history:
+            Mapping from metric name to a list of per-step metric records.
+        agent_type_counts:
+            Final observed agent counts by type.
+        step_count:
+            Number of simulated steps represented by the trajectory. This is
+            the highest observed step plus one when possible.
+        first_step:
+            First observed step, or ``None`` when there is no history.
+        last_step:
+            Last observed step, or ``None`` when there is no history.
+        snapshot_count:
+            Number of snapshots consumed.
+        metadata:
+            Additional extraction diagnostics.
+    """
+
+    population_timeseries: Mapping[int, Mapping[str, int]] = field(default_factory=dict)
+    resource_timeseries: Mapping[int, Mapping[str, float]] = field(default_factory=dict)
+    event_log: tuple[Mapping[str, Any], ...] = field(default_factory=tuple)
+    metric_history: Mapping[str, tuple[Mapping[str, Any], ...]] = field(default_factory=dict)
+    agent_type_counts: Mapping[str, int] = field(default_factory=dict)
+    step_count: int = 0
+    first_step: int | None = None
+    last_step: int | None = None
+    snapshot_count: int = 0
+    metadata: Mapping[str, Any] = field(default_factory=dict)
+
+    @property
+    def is_empty(self) -> bool:
+        """Return whether this trajectory contains no snapshots."""
+
+        return self.snapshot_count == 0
+
+    @property
+    def final_population(self) -> Mapping[str, int]:
+        """Return the final population count by agent type."""
+
+        if self.last_step is None:
+            return {}
+
+        return dict(self.population_timeseries.get(self.last_step, {}))
+
+    @property
+    def final_resources(self) -> Mapping[str, float]:
+        """Return the final resource totals by resource type."""
+
+        if self.last_step is None:
+            return {}
+
+        return dict(self.resource_timeseries.get(self.last_step, {}))
+
+    def population_delta(self) -> dict[str, int]:
+        """Return final-minus-initial population delta by agent type."""
+
+        if self.first_step is None or self.last_step is None:
+            return {}
+
+        initial = self.population_timeseries.get(self.first_step, {})
+        final = self.population_timeseries.get(self.last_step, {})
+        keys = set(initial.keys()) | set(final.keys())
+
+        return {
+            key: int(final.get(key, 0)) - int(initial.get(key, 0))
+            for key in sorted(keys)
+        }
+
+    def resource_delta(self) -> dict[str, float]:
+        """Return final-minus-initial resource amount delta by resource type."""
+
+        if self.first_step is None or self.last_step is None:
+            return {}
+
+        initial = self.resource_timeseries.get(self.first_step, {})
+        final = self.resource_timeseries.get(self.last_step, {})
+        keys = set(initial.keys()) | set(final.keys())
+
+        return {
+            key: float(final.get(key, 0.0)) - float(initial.get(key, 0.0))
+            for key in sorted(keys)
+        }
+
+    def to_dict(self) -> dict[str, Any]:
+        """Return a JSON-friendly representation of the trajectory."""
+
+        return {
+            "population_timeseries": {
+                str(step): dict(counts)
+                for step, counts in sorted(self.population_timeseries.items())
+            },
+            "resource_timeseries": {
+                str(step): dict(amounts)
+                for step, amounts in sorted(self.resource_timeseries.items())
+            },
+            "event_log": [dict(event) for event in self.event_log],
+            "metric_history": {
+                metric_name: [dict(record) for record in records]
+                for metric_name, records in self.metric_history.items()
+            },
+            "agent_type_counts": dict(self.agent_type_counts),
+            "step_count": self.step_count,
+            "first_step": self.first_step,
+            "last_step": self.last_step,
+            "snapshot_count": self.snapshot_count,
+            "population_delta": self.population_delta(),
+            "resource_delta": self.resource_delta(),
+            "metadata": copy.deepcopy(dict(self.metadata)),
+        }
+
+
+@dataclass
+class TrajectoryExtractor:
+    """Extract structured trajectory data from completed simulations."""
+
+    config: TrajectoryExtractionConfig = field(default_factory=TrajectoryExtractionConfig)
+
+    def extract(
+        self,
+        world_or_history: Any,
+        *,
+        metric_history: Mapping[str, Any] | Sequence[Any] | None = None,
+        event_log: Sequence[Any] | Mapping[str, Any] | None = None,
+    ) -> TrajectoryData:
+        """Extract trajectory data from a world-like object or raw history.
+
+        Args:
+            world_or_history:
+                Either a world-like object with ``history`` or a raw history
+                object.
+            metric_history:
+                Optional external metric history to merge with snapshot metrics.
+            event_log:
+                Optional external event log to merge with snapshot events.
+
+        Returns:
+            Normalized ``TrajectoryData``.
+        """
+
+        world = self._as_world_like(world_or_history)
+        raw_history = self._history_from_input(world_or_history)
+        snapshots = self._normalize_history(raw_history)
+
+        if not snapshots and world is not None:
+            current_snapshot = self._snapshot_from_world(world)
+            if current_snapshot:
+                snapshots = (current_snapshot,)
+
+        population_timeseries: dict[int, dict[str, int]] = {}
+        resource_timeseries: dict[int, dict[str, float]] = {}
+        normalized_event_log: list[dict[str, Any]] = []
+        normalized_metric_history: dict[str, list[dict[str, Any]]] = defaultdict(list)
+        unknown_snapshot_metadata: list[dict[str, Any]] = []
+
+        for index, raw_snapshot in enumerate(snapshots):
+            snapshot = self._mapping_from_any(raw_snapshot)
+            step = self._step_from_snapshot(snapshot, fallback=index)
+
+            population_counts = self._population_counts_from_snapshot(snapshot)
+            resource_amounts = self._resource_amounts_from_snapshot(snapshot)
+
+            population_timeseries[step] = population_counts
+            resource_timeseries[step] = resource_amounts
+
+            if self.config.include_snapshot_events:
+                self._append_snapshot_events(
+                    snapshot=snapshot,
+                    step=step,
+                    output=normalized_event_log,
+                )
+
+            if self.config.include_snapshot_metrics:
+                self._append_snapshot_metrics(
+                    snapshot=snapshot,
+                    step=step,
+                    output=normalized_metric_history,
+                )
+
+            if self.config.preserve_unknown_snapshot_fields:
+                preserved = self._preserve_unknown_snapshot_fields(snapshot, step=step)
+                if preserved:
+                    unknown_snapshot_metadata.append(preserved)
+
+        self._append_external_events(
+            event_log=event_log,
+            world=world,
+            output=normalized_event_log,
+        )
+
+        self._append_external_metrics(
+            metric_history=metric_history,
+            world=world,
+            output=normalized_metric_history,
+        )
+
+        observed_steps = sorted(population_timeseries.keys() | resource_timeseries.keys())
+        first_step = observed_steps[0] if observed_steps else None
+        last_step = observed_steps[-1] if observed_steps else None
+
+        if last_step is None:
+            step_count = int(getattr(world, "step_count", 0)) if world is not None else 0
+        else:
+            step_count = max(last_step + 1, int(getattr(world, "step_count", 0)) if world is not None else 0)
+
+        final_counts = (
+            population_timeseries.get(last_step, {})
+            if last_step is not None
+            else {}
+        )
+
+        metadata: dict[str, Any] = {
+            "source": "world" if world is not None else "history",
+            "history_length": len(snapshots),
+            "count_alive_only": self.config.count_alive_only,
+        }
+
+        if unknown_snapshot_metadata:
+            metadata["unknown_snapshot_fields"] = unknown_snapshot_metadata[
+                : self.config.max_preserved_unknown_values
+            ]
+
+        return TrajectoryData(
+            population_timeseries={
+                step: dict(counts)
+                for step, counts in sorted(population_timeseries.items())
+            },
+            resource_timeseries={
+                step: dict(amounts)
+                for step, amounts in sorted(resource_timeseries.items())
+            },
+            event_log=tuple(normalized_event_log),
+            metric_history={
+                metric_name: tuple(records)
+                for metric_name, records in sorted(normalized_metric_history.items())
+            },
+            agent_type_counts=dict(final_counts),
+            step_count=step_count,
+            first_step=first_step,
+            last_step=last_step,
+            snapshot_count=len(snapshots),
+            metadata=metadata,
+        )
+
+    def extract_from_world(self, world: WorldHistoryProtocol) -> TrajectoryData:
+        """Extract trajectory data directly from a world-like object."""
+
+        return self.extract(world)
+
+    def extract_from_history(self, history: Sequence[Any] | Mapping[str, Any]) -> TrajectoryData:
+        """Extract trajectory data directly from a raw history object."""
+
+        return self.extract(history)
+
+    @staticmethod
+    def _as_world_like(value: Any) -> Any | None:
+        """Return value when it appears to be a world-like object."""
+
+        if hasattr(value, "history"):
+            return value
+
+        return None
+
+    @staticmethod
+    def _history_from_input(value: Any) -> Any:
+        """Return raw history from a world-like object or raw history input."""
+
+        if hasattr(value, "history"):
+            return getattr(value, "history")
+
+        return value
+
+    def _normalize_history(self, raw_history: Any) -> tuple[Any, ...]:
+        """Normalize raw history into a tuple of snapshots."""
+
+        if raw_history is None:
+            return ()
+
+        if isinstance(raw_history, Mapping):
+            if "history" in raw_history:
+                return self._normalize_history(raw_history["history"])
+
+            if "snapshots" in raw_history:
+                return self._normalize_history(raw_history["snapshots"])
+
+            return (raw_history,)
+
+        if isinstance(raw_history, Sequence) and not isinstance(raw_history, (str, bytes)):
+            return tuple(raw_history)
+
+        logger.warning(
+            "TrajectoryExtractor received unsupported history type: %s",
+            raw_history.__class__.__name__,
+        )
+        return ()
+
+    def _snapshot_from_world(self, world: Any) -> dict[str, Any]:
+        """Build a single best-effort snapshot from a world-like object."""
+
+        snapshot_method = getattr(world, "snapshot", None)
+        if callable(snapshot_method):
+            try:
+                snapshot = snapshot_method()
+                if isinstance(snapshot, Mapping):
+                    return dict(snapshot)
+            except Exception:
+                logger.debug("world.snapshot() failed during trajectory extraction", exc_info=True)
+
+        to_dict = getattr(world, "to_dict", None)
+        if callable(to_dict):
+            try:
+                snapshot = to_dict()
+                if isinstance(snapshot, Mapping):
+                    return dict(snapshot)
+            except Exception:
+                logger.debug("world.to_dict() failed during trajectory extraction", exc_info=True)
+
+        snapshot: dict[str, Any] = {
+            "step": getattr(world, "step_count", 0),
+            "agents": getattr(world, "agents", {}),
+            "resources": getattr(world, "resources", {}),
+        }
+
+        metrics = getattr(world, "metrics", None)
+        if metrics is not None:
+            snapshot["metrics"] = metrics
+
+        return snapshot
+
+    def _population_counts_from_snapshot(self, snapshot: Mapping[str, Any]) -> dict[str, int]:
+        """Extract population counts by agent type from a snapshot."""
+
+        if "agent_type_counts" in snapshot and isinstance(snapshot["agent_type_counts"], Mapping):
+            return {
+                str(agent_type): int(count)
+                for agent_type, count in snapshot["agent_type_counts"].items()
+                if self._is_numeric(count)
+            }
+
+        if "population" in snapshot and isinstance(snapshot["population"], Mapping):
+            return {
+                str(agent_type): int(count)
+                for agent_type, count in snapshot["population"].items()
+                if self._is_numeric(count)
+            }
+
+        agents = self._collection_values(
+            self._first_present(
+                snapshot,
+                ("agents", "agent_states", "agent_snapshots"),
+                default=(),
+            )
+        )
+
+        counts: dict[str, int] = defaultdict(int)
+
+        for agent in agents:
+            agent_type = self._field(agent, "type", "agent_type", "kind", default="unknown")
+            alive = self._field(agent, "alive", default=True)
+
+            if self.config.count_alive_only and alive is False:
+                continue
+
+            counts[str(agent_type)] += 1
+
+        return dict(sorted(counts.items()))
+
+    def _resource_amounts_from_snapshot(self, snapshot: Mapping[str, Any]) -> dict[str, float]:
+        """Extract resource total amounts by resource type from a snapshot."""
+
+        for key in ("resource_timeseries", "resource_type_amounts", "resource_amounts"):
+            value = snapshot.get(key)
+            if isinstance(value, Mapping):
+                return {
+                    str(resource_type): float(amount)
+                    for resource_type, amount in value.items()
+                    if self._is_numeric(amount)
+                }
+
+        resources = self._collection_values(
+            self._first_present(
+                snapshot,
+                ("resources", "resource_states", "resource_snapshots"),
+                default=(),
+            )
+        )
+
+        amounts: dict[str, float] = defaultdict(float)
+
+        for resource in resources:
+            resource_type = self._field(resource, "type", "resource_type", "kind", default="unknown")
+            amount = self._field(resource, "amount", "quantity", "value", default=0.0)
+
+            if self._is_numeric(amount):
+                amounts[str(resource_type)] += float(amount)
+
+        return dict(sorted(amounts.items()))
+
+    def _append_snapshot_events(
+        self,
+        *,
+        snapshot: Mapping[str, Any],
+        step: int,
+        output: list[dict[str, Any]],
+    ) -> None:
+        """Append normalized event records from a snapshot."""
+
+        raw_events = self._first_present(
+            snapshot,
+            ("event_log", "events_processed", "events_fired", "events"),
+            default=(),
+        )
+
+        for record in self._collection_values(raw_events):
+            normalized = self._normalize_event_record(record, default_step=step)
+            if normalized:
+                output.append(normalized)
+
+    def _append_snapshot_metrics(
+        self,
+        *,
+        snapshot: Mapping[str, Any],
+        step: int,
+        output: dict[str, list[dict[str, Any]]],
+    ) -> None:
+        """Append metric records from a snapshot."""
+
+        raw_metrics = self._first_present(
+            snapshot,
+            ("metrics", "metric_values", "metric_history"),
+            default=None,
+        )
+
+        if raw_metrics is None:
+            return
+
+        self._append_metric_records(raw_metrics, default_step=step, output=output)
+
+    def _append_external_events(
+        self,
+        *,
+        event_log: Sequence[Any] | Mapping[str, Any] | None,
+        world: Any | None,
+        output: list[dict[str, Any]],
+    ) -> None:
+        """Append externally supplied or world-level event log records."""
+
+        raw_event_log = event_log
+
+        if raw_event_log is None and world is not None:
+            raw_event_log = self._first_present_from_object(
+                world,
+                ("event_log", "events_processed", "events_fired"),
+                default=None,
+            )
+
+        if raw_event_log is None:
+            return
+
+        for record in self._collection_values(raw_event_log):
+            normalized = self._normalize_event_record(record, default_step=None)
+            if normalized:
+                output.append(normalized)
+
+    def _append_external_metrics(
+        self,
+        *,
+        metric_history: Mapping[str, Any] | Sequence[Any] | None,
+        world: Any | None,
+        output: dict[str, list[dict[str, Any]]],
+    ) -> None:
+        """Append externally supplied or world-level metric history."""
+
+        raw_metric_history = metric_history
+
+        if raw_metric_history is None and world is not None:
+            raw_metric_history = self._first_present_from_object(
+                world,
+                ("metric_history", "metrics_history", "metrics"),
+                default=None,
+            )
+
+        if raw_metric_history is None:
+            return
+
+        self._append_metric_records(raw_metric_history, default_step=None, output=output)
+
+    def _append_metric_records(
+        self,
+        raw_metrics: Any,
+        *,
+        default_step: int | None,
+        output: dict[str, list[dict[str, Any]]],
+    ) -> None:
+        """Normalize and append metric records."""
+
+        if isinstance(raw_metrics, Mapping):
+            for metric_name, metric_value in raw_metrics.items():
+                if isinstance(metric_value, Sequence) and not isinstance(metric_value, (str, bytes, Mapping)):
+                    for item in metric_value:
+                        output[str(metric_name)].append(
+                            self._normalize_metric_record(
+                                item,
+                                metric_name=str(metric_name),
+                                default_step=default_step,
+                            )
+                        )
+                else:
+                    output[str(metric_name)].append(
+                        self._normalize_metric_record(
+                            metric_value,
+                            metric_name=str(metric_name),
+                            default_step=default_step,
+                        )
+                    )
+            return
+
+        if isinstance(raw_metrics, Sequence) and not isinstance(raw_metrics, (str, bytes)):
+            for item in raw_metrics:
+                item_mapping = self._mapping_from_any(item)
+
+                metric_name = str(
+                    self._first_present(
+                        item_mapping,
+                        ("name", "metric", "metric_name"),
+                        default="unknown_metric",
+                    )
+                )
+                output[metric_name].append(
+                    self._normalize_metric_record(
+                        item_mapping,
+                        metric_name=metric_name,
+                        default_step=default_step,
+                    )
+                )
+
+    def _normalize_event_record(
+        self,
+        record: Any,
+        *,
+        default_step: int | None,
+    ) -> dict[str, Any]:
+        """Normalize one event record."""
+
+        mapping = self._mapping_from_any(record)
+
+        if not mapping:
+            return {}
+
+        if default_step is not None:
+            mapping.setdefault("step", default_step)
+
+        event_name = self._first_present(
+            mapping,
+            ("name", "event", "event_name", "type"),
+            default=None,
+        )
+        if event_name is not None:
+            mapping.setdefault("name", str(event_name))
+
+        return self._json_safe_mapping(mapping)
+
+    def _normalize_metric_record(
+        self,
+        record: Any,
+        *,
+        metric_name: str,
+        default_step: int | None,
+    ) -> dict[str, Any]:
+        """Normalize one metric record."""
+
+        if isinstance(record, Mapping):
+            output = self._json_safe_mapping(record)
+        else:
+            output = {"value": self._json_safe(record)}
+
+        output.setdefault("name", metric_name)
+
+        if default_step is not None:
+            output.setdefault("step", default_step)
+
+        return output
+
+    def _preserve_unknown_snapshot_fields(
+        self,
+        snapshot: Mapping[str, Any],
+        *,
+        step: int,
+    ) -> dict[str, Any]:
+        """Preserve selected unknown fields from a snapshot for diagnostics."""
+
+        known_keys = {
+            "step",
+            "step_count",
+            "time",
+            "tick",
+            "agents",
+            "agent_states",
+            "agent_snapshots",
+            "agent_type_counts",
+            "population",
+            "resources",
+            "resource_states",
+            "resource_snapshots",
+            "resource_timeseries",
+            "resource_type_amounts",
+            "resource_amounts",
+            "events",
+            "event_log",
+            "events_processed",
+            "events_fired",
+            "metrics",
+            "metric_values",
+            "metric_history",
+        }
+
+        unknown: dict[str, Any] = {}
+
+        for key, value in snapshot.items():
+            if key in known_keys:
+                continue
+            unknown[str(key)] = self._json_safe(value)
+
+        if not unknown:
+            return {}
+
+        return {
+            "step": step,
+            "fields": unknown,
+        }
+
+    @staticmethod
+    def _step_from_snapshot(snapshot: Mapping[str, Any], *, fallback: int) -> int:
+        """Return step value from snapshot with a fallback index."""
+
+        for key in ("step", "step_count", "time", "tick"):
+            if key in snapshot and TrajectoryExtractor._is_numeric(snapshot[key]):
+                return int(snapshot[key])
+
+        return int(fallback)
+
+    @staticmethod
+    def _collection_values(value: Any) -> tuple[Any, ...]:
+        """Return collection values from mapping or sequence."""
+
+        if value is None:
+            return ()
+
+        if isinstance(value, Mapping):
+            return tuple(value.values())
+
+        if isinstance(value, Sequence) and not isinstance(value, (str, bytes)):
+            return tuple(value)
+
+        return (value,)
+
+    @staticmethod
+    def _mapping_from_any(value: Any) -> dict[str, Any]:
+        """Return a mapping representation of an arbitrary value."""
+
+        if value is None:
+            return {}
+
+        if isinstance(value, Mapping):
+            return dict(value)
+
+        for method_name in ("to_dict", "snapshot", "model_dump"):
+            method = getattr(value, method_name, None)
+            if callable(method):
+                try:
+                    result = method()
+                    if isinstance(result, Mapping):
+                        return dict(result)
+                except Exception:
+                    logger.debug(
+                        "Could not convert %s using %s()",
+                        value.__class__.__name__,
+                        method_name,
+                        exc_info=True,
+                    )
+
+        output: dict[str, Any] = {}
+
+        for attribute in ("id", "name", "type", "amount", "position", "alive", "state", "metadata"):
+            if hasattr(value, attribute):
+                output[attribute] = getattr(value, attribute)
+
+        return output
+
+    @staticmethod
+    def _field(value: Any, *names: str, default: Any = None) -> Any:
+        """Read the first present field from mapping or object."""
+
+        if isinstance(value, Mapping):
+            for name in names:
+                if name in value:
+                    return value[name]
+            return default
+
+        for name in names:
+            if hasattr(value, name):
+                return getattr(value, name)
+
+        return default
+
+    @staticmethod
+    def _first_present(mapping: Mapping[str, Any], keys: Sequence[str], *, default: Any) -> Any:
+        """Return first present mapping value."""
+
+        for key in keys:
+            if key in mapping and mapping[key] is not None:
+                return mapping[key]
+
+        return default
+
+    @staticmethod
+    def _first_present_from_object(value: Any, keys: Sequence[str], *, default: Any) -> Any:
+        """Return first present object attribute value."""
+
+        for key in keys:
+            if hasattr(value, key):
+                candidate = getattr(value, key)
+                if candidate is not None:
+                    return candidate
+
+        return default
+
+    @staticmethod
+    def _is_numeric(value: Any) -> bool:
+        """Return whether value is a finite int or float but not bool."""
+
+        if isinstance(value, bool):
+            return False
+
+        if not isinstance(value, (int, float)):
+            return False
+
+        return math.isfinite(float(value))
+
+    @classmethod
+    def _json_safe_mapping(cls, mapping: Mapping[str, Any]) -> dict[str, Any]:
+        """Return JSON-safe copy of a mapping."""
+
+        return {
+            str(key): cls._json_safe(value)
+            for key, value in mapping.items()
+        }
+
+    @classmethod
+    def _json_safe(cls, value: Any) -> Any:
+        """Return a JSON-friendly representation."""
+
+        if value is None or isinstance(value, (str, bool)):
+            return value
+
+        if isinstance(value, int) and not isinstance(value, bool):
+            return value
+
+        if isinstance(value, float):
+            return value if math.isfinite(value) else None
+
+        if isinstance(value, Mapping):
+            return {
+                str(key): cls._json_safe(item)
+                for key, item in value.items()
+            }
+
+        if isinstance(value, Sequence) and not isinstance(value, (str, bytes)):
+            return [cls._json_safe(item) for item in value]
+
+        for method_name in ("to_dict", "snapshot", "model_dump"):
+            method = getattr(value, method_name, None)
+            if callable(method):
+                try:
+                    return cls._json_safe(method())
+                except Exception:
+                    logger.debug(
+                        "Could not JSON-normalize %s using %s()",
+                        value.__class__.__name__,
+                        method_name,
+                        exc_info=True,
+                    )
+
+        if hasattr(value, "tolist") and callable(value.tolist):
+            try:
+                return cls._json_safe(value.tolist())
+            except Exception:
+                logger.debug("Could not JSON-normalize tolist() value", exc_info=True)
+
+        return repr(value)
+
+
+def extract_trajectory(
+    world_or_history: Any,
+    *,
+    metric_history: Mapping[str, Any] | Sequence[Any] | None = None,
+    event_log: Sequence[Any] | Mapping[str, Any] | None = None,
+    config: TrajectoryExtractionConfig | None = None,
+) -> TrajectoryData:
+    """Convenience function for extracting trajectory data."""
+
+    extractor = TrajectoryExtractor(config=config or TrajectoryExtractionConfig())
+    return extractor.extract(
+        world_or_history,
+        metric_history=metric_history,
+        event_log=event_log,
+    )
+
+
+__all__ = [
+    "TrajectoryData",
+    "TrajectoryExtractionConfig",
+    "TrajectoryExtractor",
+    "WorldHistoryProtocol",
+    "extract_trajectory",
+]