aethergraph 0.1.0a1__py3-none-any.whl → 0.1.0a3__py3-none-any.whl

aethergraph/services/viz/facade.py

@@ -0,0 +1,413 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Any
+
+from aethergraph.contracts.services.viz import VizEvent, VizMode
+from aethergraph.services.artifacts.facade import Artifact, ArtifactFacade
+from aethergraph.services.scope.scope import Scope
+from aethergraph.services.viz.viz_service import VizService
+
+
+@dataclass
+class VizFacade:
+    """
+    High-level facade for visualization operations within a given Scope.
+
+    - Wraps VizService and ArtifactFacade.
+    - Knows about Scope to auto-fill provenance and tenant fields.
+
+    Usage pattern in ctx.viz:
+    # Scalars
+    await ctx.viz.scalar("loss", step=iter, value=float(loss), figure_id="metrics")
+
+    # Matrix (small heatmap)
+    await ctx.viz.matrix("field_map", step=iter, matrix=field_2d, figure_id="fields")
+
+    # Image (pre-rendered PNG)
+    artifact = await ctx.artifacts.save_file(path="frame_17.png", kind="image")
+    await ctx.viz.image_from_artifact("design_shape", step=17, artifact=artifact, figure_id="design")
+    """
+
+    run_id: str
+    graph_id: str
+    node_id: str
+    tool_name: str
+    tool_version: str
+
+    viz_service: VizService
+    scope: Scope | None = None
+    artifacts: ArtifactFacade | None = None
+
+    # ------- internal helpers -------
+    def _scope_dims(self) -> dict[str, Any]:
+        if not self.scope:
+            return {}
+        return self.scope.metering_dimensions()
+
+    def _apply_scope(self, evt: VizEvent) -> VizEvent:
+        dims = self._scope_dims()
+        evt.org_id = evt.org_id or dims.get("org_id")
+        evt.user_id = evt.user_id or dims.get("user_id")
+        evt.client_id = evt.client_id or dims.get("client_id")
+        evt.app_id = evt.app_id or dims.get("app_id")
+        evt.session_id = evt.session_id or dims.get("session_id")
+        return evt
+
+    # ------- public API -------
+    async def scalar(
+        self,
+        track_id: str,
+        *,
+        step: int,
+        value: float,
+        figure_id: str | None = None,
+        mode: VizMode = "append",
+        meta: dict[str, Any] | None = None,
+        tags: list[str] | None = None,
+    ) -> None:
+        """
+        Record a single scalar value for visualization in the Aethergraph UI.
+
+        This method standardizes the event format, auto-fills provenance fields,
+        and dispatches the scalar data to the configured storage backend.
+
+        Examples:
+            Basic usage to log a loss metric:
+            ```python
+            await context.viz().scalar("loss", step=iteration, value=loss)
+            ```
+
+            Logging a scalar with extra metadata and custom tags:
+            ```python
+            await context.viz().scalar(
+                "accuracy",
+                step=42,
+                value=0.98,
+                figure_id="metrics",
+                meta={"model": "resnet"},
+                tags=["experiment:baseline"]
+            )
+            ```
+
+        Args:
+            track_id: Unique identifier for the scalar track (e.g., "loss").
+            step: Integer step or iteration number for the data point.
+            value: The scalar value to record (float).
+            figure_id: Optional figure grouping for UI display.
+            mode: Storage mode, typically "append".
+            meta: Optional dictionary of extra metadata.
+            tags: Optional list of string labels. The tag "type:scalar" is automatically appended.
+
+        Returns:
+            None. The event is persisted for later visualization.
+        """
+        evt = VizEvent(
+            run_id=self.run_id,
+            graph_id=self.graph_id,
+            node_id=self.node_id,
+            tool_name=self.tool_name,
+            tool_version=self.tool_version,
+            track_id=track_id,
+            figure_id=figure_id,
+            viz_kind="scalar",
+            step=step,
+            mode=mode,
+            value=float(value),
+            meta=meta,
+            tags=(tags or []) + ["type:scalar"],
+        )
+        evt = self._apply_scope(evt)
+        await self.viz_service.append(evt)
+
+    async def vector(
+        self,
+        track_id: str,
+        *,
+        step: int,
+        values: Sequence[float],
+        figure_id: str | None = None,
+        mode: VizMode = "append",
+        meta: dict[str, Any] | None = None,
+        tags: list[str] | None = None,
+    ) -> None:
+        """
+        Record a single vector (1D array) for visualization in the Aethergraph UI.
+
+        This method standardizes the event format, auto-fills provenance fields,
+        and dispatches the vector data to the configured storage backend.
+
+        Examples:
+            Basic usage to log a vector:
+            ```python
+            await context.viz().vector("embedding", step=iteration, values=[0.1, 0.2, 0.3])
+            ```
+
+            Logging a vector with extra metadata and custom tags:
+            ```python
+            await context.viz().vector(
+                "features",
+                step=42,
+                values=[1.0, 2.5, 3.7],
+                figure_id="feature_tracks",
+                meta={"source": "encoder"},
+                tags=["experiment:baseline"]
+            )
+            ```
+
+        Args:
+            track_id: Unique identifier for the vector track (e.g., "embedding").
+            step: Integer step or iteration number for the data point.
+            values: Sequence of float values representing the vector.
+            figure_id: Optional figure grouping for UI display.
+            mode: Storage mode, typically "append".
+            meta: Optional dictionary of extra metadata.
+            tags: Optional list of string labels. The tag "type:vector" is automatically appended.
+
+        Returns:
+            None. The event is persisted for later visualization.
+        """
+        evt = VizEvent(
+            run_id=self.run_id,
+            graph_id=self.graph_id,
+            node_id=self.node_id,
+            tool_name=self.tool_name,
+            tool_version=self.tool_version,
+            track_id=track_id,
+            figure_id=figure_id,
+            viz_kind="vector",
+            step=step,
+            mode=mode,
+            vector=[float(v) for v in values],
+            meta=meta,
+            tags=(tags or []) + ["type:vector"],
+        )
+        evt = self._apply_scope(evt)
+        await self.viz_service.append(evt)
+
+    async def matrix(
+        self,
+        track_id: str,
+        *,
+        step: int,
+        matrix: Sequence[Sequence[float]],
+        figure_id: str | None = None,
+        mode: VizMode = "append",
+        meta: dict[str, Any] | None = None,
+        tags: list[str] | None = None,
+    ) -> None:
+        """
+        Record a single matrix (2D array) for visualization in the Aethergraph UI.
+
+        This method standardizes the event format, auto-fills provenance fields,
+        and dispatches the matrix data to the configured storage backend.
+
+        Examples:
+            Basic usage to log a matrix:
+            ```python
+            await context.viz().matrix("confusion", step=iteration, matrix=[[1, 2], [3, 4]])
+            ```
+
+            Logging a matrix with extra metadata and custom tags:
+            ```python
+            await context.viz().matrix(
+                "heatmap",
+                step=42,
+                matrix=[[0.1, 0.2], [0.3, 0.4]],
+                figure_id="metrics",
+                meta={"source": "model"},
+                tags=["experiment:baseline"]
+            )
+            ```
+
+        Args:
+            track_id: Unique identifier for the matrix track (e.g., "confusion").
+            step: Integer step or iteration number for the data point.
+            matrix: Sequence of sequences of float values representing the 2D matrix.
+            figure_id: Optional figure grouping for UI display.
+            mode: Storage mode, typically "append".
+            meta: Optional dictionary of extra metadata.
+            tags: Optional list of string labels. The tag "matrix" is automatically appended.
+
+        Returns:
+            None. The event is persisted for later visualization.
+        """
+        # Convert to plain list[list[float]]
+        m = [[float(x) for x in row] for row in matrix]
+        evt = VizEvent(
+            run_id=self.run_id,
+            graph_id=self.graph_id,
+            node_id=self.node_id,
+            tool_name=self.tool_name,
+            tool_version=self.tool_version,
+            track_id=track_id,
+            figure_id=figure_id,
+            viz_kind="matrix",
+            step=step,
+            mode=mode,
+            matrix=m,
+            meta=meta,
+            tags=(tags or []) + ["matrix"],
+        )
+        evt = self._apply_scope(evt)
+        await self.viz_service.append(evt)
+
+    async def image_from_artifact(
+        self,
+        track_id: str,
+        *,
+        step: int,
+        artifact: Artifact,
+        figure_id: str | None = None,
+        mode: VizMode = "append",
+        meta: dict[str, Any] | None = None,
+        tags: list[str] | None = None,
+    ) -> None:
+        """
+        Record a reference to an existing image Artifact for visualization in the Aethergraph UI.
+
+        This method standardizes the event format, auto-fills provenance fields,
+        and dispatches the image reference to the configured storage backend.
+
+        Examples:
+            Basic usage to log an image artifact:
+            ```python
+            await context.viz().image_from_artifact(
+                "design_shape",
+                step=17,
+                artifact=artifact,
+                figure_id="design"
+            )
+            ```
+
+            Logging an image with extra metadata and custom tags:
+            ```python
+            await context.viz().image_from_artifact(
+                "output_frame",
+                step=42,
+                artifact=artifact,
+                figure_id="frames",
+                meta={"source": "simulation"},
+                tags=["experiment:baseline"]
+            )
+            ```
+
+        Args:
+            track_id: Unique identifier for the image track (e.g., "design_shape").
+            step: Integer step or iteration number for the data point.
+            artifact: The Artifact object referencing the stored image.
+            figure_id: Optional figure grouping for UI display.
+            mode: Storage mode, typically "append".
+            meta: Optional dictionary of extra metadata.
+            tags: Optional list of string labels. The tag "image" is automatically appended.
+
+        Returns:
+            None. The event is persisted for later visualization.
+        """
+        evt = VizEvent(
+            run_id=self.run_id,
+            graph_id=self.graph_id,
+            node_id=self.node_id,
+            tool_name=self.tool_name,
+            tool_version=self.tool_version,
+            track_id=track_id,
+            figure_id=figure_id,
+            viz_kind="image",
+            step=step,
+            mode=mode,
+            artifact_id=artifact.artifact_id,
+            meta=meta,
+            tags=(tags or []) + ["image"],
+        )
+        evt = self._apply_scope(evt)
+        await self.viz_service.append(evt)
+
+    async def image_from_bytes(
+        self,
+        track_id: str,
+        *,
+        step: int,
+        data: bytes,
+        mime: str = "image/png",
+        kind: str = "image",
+        figure_id: str | None = None,
+        mode: VizMode = "append",
+        labels: dict[str, Any] | None = None,
+        meta: dict[str, Any] | None = None,
+        tags: list[str] | None = None,
+    ) -> Artifact:
+        """
+        Save image bytes as an Artifact and log a visualization event.
+
+        This convenience method is accessed via `context.viz().image_from_bytes(...)` and is used by the Aethergraph UI to persist image data to storage. It stores the image as an Artifact using the configured ArtifactFacade, then logs a visualization event referencing the saved artifact.
+
+        Examples:
+            Saving a PNG image to the current visualization track:
+            ```python
+            await context.viz().image_from_bytes(
+                track_id="experiment-123",
+                step=42,
+                data=image_bytes,
+                mime="image/png",
+                labels={"type": "output", "stage": "inference"},
+                tags=["result", "png"]
+            )
+            ```
+
+            Saving an image with custom metadata and figure association:
+            ```python
+            await context.viz().image_from_bytes(
+                track_id="demo-track",
+                step=7,
+                data=img_bytes,
+                figure_id="fig-1",
+                meta={"caption": "Sample output"},
+                mode="replace"
+            )
+            ```
+
+        Args:
+            track_id: The identifier for the visualization track to associate with the image.
+            step: The step or index within the track for this image.
+            data: Raw image bytes to be saved.
+            mime: The MIME type of the image (default: "image/png").
+            kind: The artifact kind (default: "image").
+            figure_id: Optional identifier for the figure this image belongs to.
+            mode: Visualization mode, e.g., "append" or "replace".
+            labels: Optional dictionary of labels to attach to the artifact.
+            meta: Optional dictionary of metadata for the visualization event.
+            tags: Optional list of string tags for categorization.
+
+        Returns:
+            Artifact: The persisted `Artifact` object representing the saved image.
+
+        Notes:
+            - This method requires that `self.artifacts` is set to an `ArtifactFacade` instance.
+            - The saved artifact is automatically linked to the visualization event.
+            - To change the name of the saved artifact in UI, use `labels` to set a "filename" label.
+        """
+        if not self.artifacts:
+            raise RuntimeError("VizFacade.image_from_bytes requires an ArtifactFacade")
+
+        # Save artifact using writer() so we get proper metering + labels
+
+        # Use ArtifactFacade.writer to store the image
+        async with self.artifacts.writer(kind=kind, planned_ext=".png") as w:
+            w.write(data)
+            if labels:
+                w.add_labels(labels)
+        art = self.artifacts.last_artifact
+        if not art:
+            raise RuntimeError("Artifact writer did not produce an artifact")
+
+        await self.image_from_artifact(
+            track_id=track_id,
+            step=step,
+            artifact=art,
+            figure_id=figure_id,
+            mode=mode,
+            meta=meta,
+            tags=tags,
+        )
+        return art

aethergraph/services/viz/viz_service.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+from dataclasses import asdict
+from datetime import datetime, timezone
+from typing import Any
+
+from aethergraph.contracts.services.viz import VizEvent, VizKind
+from aethergraph.contracts.storage.event_log import EventLog
+
+
+class VizService:
+    """
+    Low-level service to append/query visualization events.
+
+    - Uses EventLog as the underlying storage.
+    - Does NOT know about NodeContext or Scope; that's the Facade's job.
+    """
+
+    def __init__(self, event_log: EventLog):
+        self._log = event_log
+
+    async def append(self, evt: VizEvent) -> None:
+        now = datetime.now(timezone.utc).isoformat()
+        if not evt.created_at:
+            evt.created_at = now
+
+        payload = asdict(evt)
+        await self._log.append(
+            {
+                "kind": "viz",
+                "scope_id": evt.run_id,
+                "ts": evt.created_at,
+                "data": payload,
+                "tags": (evt.tags or []) + [f"track:{evt.track_id}"],
+            }
+        )
+
+    async def query_run(
+        self,
+        run_id: str,
+        *,
+        kinds: list[VizKind] | None = None,
+        since: datetime | None = None,
+        until: datetime | None = None,
+        limit: int | None = None,
+        offset: int = 0,
+    ) -> list[dict[str, Any]]:
+        """
+        Raw fetch of viz events for a given run.
+        Returns raw event dicts as stored in EventLog.
+        The API layer can normalize/group them for the frontend.
+        """
+        rows = await self._log.query(
+            scope_id=run_id,
+            since=since,
+            until=until,
+            kinds=["viz"],
+            limit=limit,
+            offset=offset,
+        )
+        # Optionally filter by viz_kind inside data
+        if kinds:
+            out: list[dict[str, Any]] = []
+            for r in rows:
+                data = r.get("data") or {}
+                if data.get("viz_kind") in kinds:
+                    out.append(r)
+            return out
+        return rows

aethergraph/storage/artifacts/artifact_index_jsonl.py

@@ -0,0 +1,180 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+import threading
+from typing import Any, Literal
+
+from aethergraph.contracts.services.artifacts import Artifact
+from aethergraph.contracts.storage.artifact_index import AsyncArtifactIndex
+
+
+class JsonlArtifactIndexSync:
+    """
+    Simple JSONL-based artifact index for small/medium scale.
+    Not suitable for millions of artifacts due to linear scans.
+    """
+
+    def __init__(self, path: str, occurrences_path: str | None = None):
+        self.path = path
+        self.occ_path = occurrences_path or (os.path.splitext(path)[0] + "_occurrences.jsonl")
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+
+        self._by_id: dict[str, dict[str, Any]] = {}
+        self._lock = threading.Lock()
+
+        if os.path.exists(self.path):
+            with open(self.path, encoding="utf-8") as f:
+                for line in f:
+                    if not line.strip():
+                        continue
+                    rec = json.loads(line)
+                    self._by_id[rec["artifact_id"]] = rec
+
+    # -------- core operations --------
+
+    def upsert(self, a: Artifact) -> None:
+        with self._lock:
+            rec = a.to_dict()
+            self._by_id[a.artifact_id] = rec
+            with open(self.path, "a", encoding="utf-8") as f:
+                f.write(json.dumps(rec) + "\n")
+
+    def list_for_run(self, run_id: str) -> list[Artifact]:
+        return [Artifact(**r) for r in self._by_id.values() if r.get("run_id") == run_id]
+
+    def search(
+        self,
+        *,
+        kind: str | None = None,
+        labels: dict[str, Any] | None = None,
+        metric: str | None = None,
+        mode: Literal["max", "min"] | None = None,
+        limit: int | None = None,
+        offset: int = 0,
+    ) -> list[Artifact]:
+        # NOTE: JSONL index keeps all artifacts in memory (_by_id.values()) and
+        # performs filtering / sorting in Python, then applies offset + limit.
+        # This is intended for small/medium local installs and tests only.
+        rows = list(self._by_id.values())
+
+        if kind:
+            rows = [r for r in rows if r.get("kind") == kind]
+
+        # Treat tenant keys as top-level fields, not labels
+        TENANT_KEYS = {
+            "org_id",
+            "user_id",
+            "client_id",
+            "app_id",
+            "session_id",
+            "run_id",
+        }
+
+        if labels:
+            for k, v in labels.items():
+                if k in TENANT_KEYS:
+                    # Match against top-level JSON fields
+                    rows = [r for r in rows if r.get(k) == v]
+                    continue
+
+                # Normal label filters
+                if isinstance(v, list):
+                    rows = [
+                        r
+                        for r in rows
+                        if isinstance(r.get("labels", {}).get(k), list)
+                        and set(v).issubset(set(r["labels"][k]))
+                    ]
+                else:
+                    rows = [r for r in rows if r.get("labels", {}).get(k) == v]
+
+        if metric and mode:
+            rows = [r for r in rows if metric in r.get("metrics", {})]
+            rows.sort(
+                key=lambda r: r["metrics"][metric],
+                reverse=(mode == "max"),
+            )
+
+        if offset > 0:
+            rows = rows[offset:]
+
+        if limit is not None:
+            rows = rows[:limit]
+
+        return [Artifact(**r) for r in rows]
+
+    def best(
+        self,
+        *,
+        kind: str,
+        metric: str,
+        mode: Literal["max", "min"],
+        filters: dict[str, Any] | None = None,
+    ) -> Artifact | None:
+        rows = self.search(
+            kind=kind,
+            labels=filters,
+            metric=metric,
+            mode=mode,
+            limit=1,
+        )
+        return rows[0] if rows else None
+
+    def pin(self, artifact_id: str, pinned: bool = True) -> None:
+        with self._lock:
+            if artifact_id not in self._by_id:
+                return
+            rec = self._by_id[artifact_id]
+            rec["pinned"] = bool(pinned)
+            self._by_id[artifact_id] = rec
+            with open(self.path, "a", encoding="utf-8") as f:
+                f.write(json.dumps(rec) + "\n")
+
+    def record_occurrence(self, a: Artifact, extra_labels: dict | None = None) -> None:
+        row = {
+            "artifact_id": a.artifact_id,
+            "run_id": a.run_id,
+            "graph_id": a.graph_id,
+            "node_id": a.node_id,
+            "tool_name": a.tool_name,
+            "tool_version": a.tool_version,
+            "created_at": a.created_at,
+            "labels": {**(a.labels or {}), **(extra_labels or {})},
+        }
+        with open(self.occ_path, "a", encoding="utf-8") as f:
+            f.write(json.dumps(row) + "\n")
+
+    def get(self, artifact_id: str) -> Artifact | None:
+        if artifact_id in self._by_id:
+            return Artifact(**self._by_id[artifact_id])
+        return None
+
+
+class JsonlArtifactIndex(AsyncArtifactIndex):
+    """Async wrapper for JsonlArtifactIndexSync using asyncio.to_thread."""
+
+    def __init__(self, path: str, occurrences_path: str | None = None):
+        self._sync = JsonlArtifactIndexSync(path, occurrences_path)
+
+    async def upsert(self, a: Artifact) -> None:
+        await asyncio.to_thread(self._sync.upsert, a)
+
+    async def list_for_run(self, run_id: str) -> list[Artifact]:
+        return await asyncio.to_thread(self._sync.list_for_run, run_id)
+
+    async def search(self, **kwargs) -> list[Artifact]:
+        return await asyncio.to_thread(self._sync.search, **kwargs)
+
+    async def best(self, **kwargs) -> Artifact | None:
+        return await asyncio.to_thread(self._sync.best, **kwargs)
+
+    async def pin(self, artifact_id: str, pinned: bool = True) -> None:
+        await asyncio.to_thread(self._sync.pin, artifact_id, pinned)
+
+    async def record_occurrence(self, a: Artifact, extra_labels: dict | None = None) -> None:
+        await asyncio.to_thread(self._sync.record_occurrence, a, extra_labels)
+
+    async def get(self, artifact_id: str) -> Artifact | None:
+        return await asyncio.to_thread(self._sync.get, artifact_id)