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

aethergraph/contracts/services/sessions.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+from datetime import datetime
+from typing import Protocol
+
+from aethergraph.api.v1.schemas import Session
+from aethergraph.core.runtime.run_types import SessionKind
+
+
+class SessionStore(Protocol):
+    async def create(
+        self,
+        *,
+        kind: SessionKind,
+        user_id: str | None = None,
+        org_id: str | None = None,
+        title: str | None = None,
+        source: str = "webui",
+        external_ref: str | None = None,
+    ) -> Session:
+        """
+        Create a new session and return it.
+        """
+
+    async def get(self, session_id: str) -> Session | None:
+        """
+        Get a session by its ID, or None if not found.
+        """
+
+    async def list_for_user(
+        self,
+        *,
+        user_id: str | None,
+        org_id: str | None = None,
+        kind: SessionKind | None = None,
+        limit: int = 50,
+        offset: int = 0,
+    ) -> Sequence[Session]:
+        """
+        List sessions for a specific user, optionally filtered by kind.
+        """
+
+    async def touch(
+        self,
+        session_id: str,
+        *,
+        updated_at: datetime | None = None,
+    ) -> None:
+        """
+        Update session's updated_at (e.g., when new message/run occurs).
+        No-op if session doesn't exist.
+        """
+
+    async def update(
+        self,
+        session_id: str,
+        *,
+        title: str | None = None,
+        external_ref: str | None = None,
+    ) -> Session | None:
+        """
+        Update session metadata, returning the updated session.
+        No-op if session doesn't exist (returns None).
+        """
+
+    async def delete(self, session_id: str) -> None:
+        """
+        Delete a session by its ID.
+        No-op if session doesn't exist.
+        """
+
+    async def record_artifact(
+        self,
+        session_id: str,
+        *,
+        created_at: datetime | None = None,
+    ) -> None:
+        """
+        Update artifact-related stats for a session:
+
+          - increment artifact_count
+          - update last_artifact_at
+
+        No-op if session_id does not exist.
+        """
+        ...

aethergraph/contracts/services/state_stores.py

@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from datetime import datetime
 from typing import Any, Protocol
 
 
@@ -13,6 +14,8 @@ class GraphSnapshot:
     created_at: float  # epoch seconds
     spec_hash: str  # detect spec drift
     state: dict[str, Any]  # JSON-serializable TaskGraphState
+    started_at: datetime | None = None
+    finished_at: datetime | None = None
 
 
 @dataclass

aethergraph/contracts/services/viz.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Literal
+
+VizKind = Literal["scalar", "vector", "matrix", "image"]
+VizMode = Literal["append", "replace"]
+
+
+@dataclass
+class VizEvent:
+    # Provenance
+    run_id: str
+    graph_id: str
+    node_id: str
+    tool_name: str
+    tool_version: str
+
+    # Visualization fields
+    track_id: str  # unique id for the trace (e.g., "loss", "accuracy")
+    figure_id: str  # optional figure id for grouping traces, e.g. "metrics_panel"
+    viz_kind: VizKind
+    step: int  # iteration or step number
+    mode: VizMode = "append"  # append or replace
+
+    # Tenant-ish fields
+    org_id: str | None = None
+    user_id: str | None = None
+    client_id: str | None = None
+    app_id: str | None = None
+    session_id: str | None = None
+
+    # Payload
+    value: float | None = None  # for scalar
+    vector: list[float] | None = None  # for vector
+    matrix: list[list[float]] | None = None  # for matrix
+    artifact_id: str | None = None  # for image or other artifact-based viz
+
+    # Optional metadata
+    meta: dict[str, Any] | None = None  # {"label": "Training Loss", "color": "blue", ...}
+    tags: list[str] | None = None  # arbitrary tags for filtering or grouping
+
+    # Timestamp
+    created_at: str | None = None  # ISO 8601 timestamp

aethergraph/contracts/storage/artifact_index.py

@@ -0,0 +1,88 @@
+from typing import Any, Literal, Protocol
+
+from aethergraph.contracts.services.artifacts import Artifact
+
+"""
+Artifact index interface for storing and retrieving artifact metadata.
+This is a special index used for tracking artifacts generated during runs.
+
+Typical implementations include:
+- FileSystemArtifactIndex: File system-based artifact index for durable storage
+- DatabaseArtifactIndex: (future) Database-backed artifact index for scalable storage and querying
+
+Note Artifact index is a specialized index for artifacts, different from general document or blob stores. 
+"""
+
+
+class AsyncArtifactIndex(Protocol):
+    """Backend-agnostic index for artifact metadata & occurrences."""
+
+    async def upsert(self, a: Artifact) -> None:
+        """Insert or update a single artifact record."""
+        ...
+
+    async def list_for_run(self, run_id: str) -> list[Artifact]:
+        """Return all artifacts for a given run_id."""
+        ...
+
+    async 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]:
+        """
+        Generic search:
+          - kind: filter by artifact.kind
+          - labels: exact-match filter on labels[k] == v
+          - metric: if provided with mode, order by metrics[metric]
+          - mode: "max" or "min" for metric
+          - limit: optional max number of results
+          - offset: pagination offset, i.e. skip this many rows before returning results
+        """
+        ...
+
+    async def best(
+        self,
+        *,
+        kind: str,
+        metric: str,
+        mode: Literal["max", "min"],
+        filters: dict[str, Any] | None = None,
+    ) -> Artifact | None:
+        """Return the single best artifact for metric under optional label filters."""
+        ...
+
+    async def pin(self, artifact_id: str, pinned: bool = True) -> None:
+        """Mark/unmark an artifact as pinned."""
+        ...
+
+    async def record_occurrence(
+        self,
+        a: Artifact,
+        extra_labels: dict | None = None,
+    ) -> None:
+        """
+        Append-only lineage log: "this artifact was used/created here".
+        """
+        ...
+
+    async def get(self, artifact_id: str) -> Artifact | None:
+        """Get artifact by ID."""
+        ...
+
+    # TODO: add cursor-based pagination for listing/searching large sets
+    # e.g.
+    # async def search_paginated(
+    #     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,
+    #     cursor: str | None = None,
+    # ) -> list[Artifact]:

aethergraph/contracts/storage/artifact_store.py

@@ -0,0 +1,99 @@
+from contextlib import AbstractAsyncContextManager
+from typing import Any, Protocol
+
+from aethergraph.contracts.services.artifacts import Artifact
+
+
+class AsyncArtifactStore(Protocol):
+    @property
+    def base_uri(self) -> str: ...
+
+    # ---------- save / ingest ----------
+    async def save_file(
+        self,
+        *,
+        path: str,
+        kind: str,
+        run_id: str,
+        graph_id: str,
+        node_id: str,
+        tool_name: str,
+        tool_version: str,
+        suggested_uri: str | None = None,
+        pin: bool = False,
+        labels: dict | None = None,
+        metrics: dict | None = None,
+        preview_uri: str | None = None,
+    ) -> Artifact: ...
+
+    async def open_writer(
+        self,
+        *,
+        kind: str,
+        run_id: str,
+        graph_id: str,
+        node_id: str,
+        tool_name: str,
+        tool_version: str,
+        planned_ext: str | None = None,
+        pin: bool = False,
+    ) -> AbstractAsyncContextManager[Any]: ...
+
+    async def plan_staging_path(self, planned_ext: str = "") -> str: ...
+    async def ingest_staged_file(
+        self,
+        *,
+        staged_path: str,
+        kind: str,
+        run_id: str,
+        graph_id: str,
+        node_id: str,
+        tool_name: str,
+        tool_version: str,
+        pin: bool = False,
+        labels: dict | None = None,
+        metrics: dict | None = None,
+        preview_uri: str | None = None,
+        suggested_uri: str | None = None,
+    ) -> Artifact: ...
+
+    async def plan_staging_dir(self, suffix: str = "") -> str: ...
+    async def ingest_directory(
+        self,
+        *,
+        staged_dir: str,
+        kind: str,
+        run_id: str,
+        graph_id: str,
+        node_id: str,
+        tool_name: str,
+        tool_version: str,
+        include: list[str] | None = None,
+        exclude: list[str] | None = None,
+        index_children: bool = False,
+        pin: bool = False,
+        labels: dict | None = None,
+        metrics: dict | None = None,
+        suggested_uri: str | None = None,
+        archive: bool = False,
+        archive_name: str = "bundle.tar.gz",
+        cleanup: bool = True,
+        store: str | None = None,
+    ) -> Artifact: ...
+
+    # ---------- load ----------
+    async def load_bytes(self, uri: str) -> bytes: ...
+    async def load_text(
+        self,
+        uri: str,
+        *,
+        encoding: str = "utf-8",
+        errors: str = "strict",
+    ) -> str: ...
+
+    async def load_artifact(self, uri: str) -> Any: ...
+    async def load_artifact_bytes(self, uri: str) -> bytes: ...
+    async def load_artifact_dir(self, uri: str) -> str: ...
+
+    # ---------- housekeeping ----------
+    async def cleanup_tmp(self, max_age_hours: int = 24) -> None: ...

aethergraph/contracts/storage/async_kv.py

@@ -0,0 +1,34 @@
+from typing import Any, Protocol
+
+"""
+Used for defining the AsyncKV protocol interface.
+
+Typical implementations include:
+- EphemeralKV: In-memory, transient key-value store.
+- SQLiteKV: Durable key-value store backed by SQLite.
+- LayeredKV: Combines a fast ephemeral cache with a durable backend.
+- RedisKV: (future) Cloud-based key-value store using Redis.
+- Factory function to create KV instances based on environment configuration.
+
+It is used in various parts of the system for transient and durable storage needs.
+- context.kv() for general KV storage.
+- memory hotlog implementation with KV support 
+"""
+
+
+class AsyncKV(Protocol):
+    async def get(self, key: str, default: Any = None) -> Any: ...
+    async def set(self, key: str, value: Any, *, ttl_s: int | None = None) -> None: ...
+    async def delete(self, key: str) -> None: ...
+
+    # Optional
+    async def mget(self, keys: list[str]) -> list[Any]: ...  # multiple get
+    async def mset(
+        self, kv: dict[str, Any], *, ttl_s: int | None = None
+    ) -> None: ...  # multiple set
+    async def expire(self, key: str, ttl_s: int) -> None: ...
+    async def purge_expired(self, limit: int = 1000) -> int: ...  # return number purged
+
+    # Optional: if implemented, allows scanning for cleanup and debugging
+    # Should return all keys starting with "prefix"
+    async def scan_keys(self, prefix: str) -> list[str]: ...

aethergraph/contracts/storage/blob_store.py

@@ -0,0 +1,50 @@
+from typing import Protocol
+
+"""
+Blob store interface for storing and retrieving binary large objects (blobs).
+Typical implementations include:
+- FSBlobStore: File system-based blob store for persistent storage
+- S3BlobStore: Amazon S3-based blob store for cloud storage
+
+This is usually used for storing large binary data such as images, audio, or other media files.
+For example:
+- Artifact Store for saving generated images or files
+"""
+
+
+class BlobStore(Protocol):
+    @property
+    def base_uri(self) -> str:  # e.g. file:///..., s3://bucket/prefix
+        ...
+
+    async def put_bytes(
+        self,
+        data: bytes,
+        *,
+        key: str | None = None,
+        ext: str | None = None,
+        mime: str | None = None,
+    ) -> str:
+        """Store bytes under an optional key; return full blob URI."""
+        ...
+
+    async def put_file(
+        self,
+        path: str,
+        *,
+        key: str | None = None,
+        mime: str | None = None,
+        keep_source: bool = False,  # whether to keep the source file after storing, only relevant for FSBlobStore
+    ) -> str:
+        """Store a local file; return full blob URI."""
+        ...
+
+    async def load_bytes(self, uri: str) -> bytes: ...
+
+    async def load_text(
+        self,
+        uri: str,
+        *,
+        encoding: str = "utf-8",
+        errors: str = "strict",
+    ) -> str: ...

aethergraph/contracts/storage/doc_store.py

@@ -0,0 +1,35 @@
+from typing import Any, Protocol
+
+"""
+Document store interface for storing and retrieving JSON-like documents.
+
+Typical implementations include:
+- InMemoryDocStore: Transient, in-memory document store for testing or ephemeral use cases
+- FileSystemDocStore: File system-based document store for durable storage
+- DatabaseDocStore: (future) Database-backed document store for scalable storage and querying
+
+It is used in various parts of the system for storing structured documents.
+- memory persistence saving summary JSON documents
+- continuation storage for saving intermediate results and tokens
+- graph state store for saving state snapshots
+"""
+
+
+class DocStore(Protocol):
+    """
+    Generic doc_id → JSON dict store.
+
+    NOTE:
+      - This is intentionally low-level and schema-agnostic.
+      - Higher-level stores (RunStore, EventLog, ArtifactIndex, etc.) are responsible
+        for applying domain-specific filtering, sorting, and pagination.
+      - DocStore.list() may scan all docs; it is not intended as a scalable
+        query interface for large, structured datasets.
+    """
+
+    async def put(self, doc_id: str, doc: dict[str, Any]) -> None: ...
+    async def get(self, doc_id: str) -> dict[str, Any] | None: ...
+    async def delete(self, doc_id: str) -> None: ...
+
+    # Optional
+    async def list(self) -> list[str]: ...

aethergraph/contracts/storage/event_log.py

@@ -0,0 +1,31 @@
+from datetime import datetime
+from typing import Protocol
+
+"""
+Event log interface for appending and querying events.
+
+Typical implementations include:
+- InMemoryEventLog: Transient, in-memory event log for testing or ephemeral use cases
+- FSPersistenceEventLog: File system-based event log for durable storage
+- DatabaseEventLog: (future) Database-backed event log for scalable storage and querying
+
+It is used in various parts of the system for logging events with metadata.
+- memory persistent implementation for saving events durably
+- graph state store for appending state change events
+"""
+
+
+class EventLog(Protocol):
+    async def append(self, evt: dict) -> None: ...
+
+    async def query(
+        self,
+        *,
+        scope_id: str | None = None,  # filter by scope ID, e.g., run ID, memory ID
+        since: datetime | None = None,  # filter events after this time
+        until: datetime | None = None,  # filter events before this time
+        kinds: list[str] | None = None,  # filter by event kinds
+        limit: int | None = None,  # max number of events to return
+        tags: list[str] | None = None,  # filter by tags
+        offset: int = 0,  # pagination offset
+    ) -> list[dict]: ...

aethergraph/contracts/storage/vector_index.py

@@ -0,0 +1,48 @@
+from typing import Any, Protocol
+
+"""
+Vector index interface for storing and retrieving vector embeddings.
+
+It can be used in rag services or any system that requires vector similarity search.
+"""
+
+
+class VectorIndex(Protocol):
+    async def add(
+        self,
+        corpus_id: str,
+        chunk_ids: list[str],
+        vectors: list[list[float]],
+        metas: list[dict[str, Any]],
+    ) -> None:
+        """
+        Insert or upsert vectors into a corpus.
+
+        - corpus_id: logical collection name
+        - chunk_ids: user IDs for each vector
+        - vectors: len == len(chunk_ids), each a dense float vector
+        - metas: arbitrary metadata (e.g. {"doc_id": ..., "offset": ...})
+        """
+
+    async def delete(
+        self,
+        corpus_id: str,
+        chunk_ids: list[str] | None = None,
+    ) -> None:
+        """
+        Delete entire corpus (chunk_ids=None) or specific chunks.
+        """
+
+    async def search(
+        self,
+        corpus_id: str,
+        query_vec: list[float],
+        k: int,
+    ) -> list[dict[str, Any]]: ...
+
+    # Each dict MUST look like:
+    # {"chunk_id": str, "score": float, "meta": dict[str, Any]}
+
+    # Optional
+    async def list_corpora(self) -> list[str]: ...
+    async def list_chunks(self, corpus_id: str) -> list[str]: ...

aethergraph/core/execution/forward_scheduler.py

@@ -200,6 +200,9 @@ class ForwardScheduler(BaseScheduler):
         # logger
         self.logger = logger
 
+        # termination flag
+        self._cancelled = False
+
     def bind_loop(self, loop: asyncio.AbstractEventLoop | None = None):
         """Bind an event loop to this scheduler (for cross-thread resume calls)."""
         self.loop = loop or asyncio.get_running_loop()
@@ -322,6 +325,12 @@ class ForwardScheduler(BaseScheduler):
                 if not ctrl.done():
                     ctrl.cancel()
 
+        if self._cancelled:
+            # propagate an explicit cancellation upwards
+            raise asyncio.CancelledError(
+                f"ForwardScheduler for run_id={self.env.run_id} was terminated"
+            )
+
     async def run_from(self, node_ids: list[str]):
         """Run starting from specific nodes (e.g. after external event)."""
         for nid in node_ids:
@@ -340,6 +349,8 @@ class ForwardScheduler(BaseScheduler):
     async def terminate(self):
         """Terminate execution; running tasks will complete but no new tasks will be started."""
         self._terminated = True
+        self._cancelled = True
+
         # cancel backoff tasks
         for task in self._backoff_tasks.values():
             task.cancel()
@@ -700,10 +711,10 @@ class ForwardScheduler(BaseScheduler):
 
             except NotImplementedError:
                 # subgraph logic not handled here; escalate to orchestrator
-                await node.set_status(NodeStatus.FAILED)
+                await self.graph.set_node_status(node_id, NodeStatus.FAILED)
             except asyncio.CancelledError:
                 # task cancelled (e.g. on terminate);
-                await node.set_status(NodeStatus.FAILED)
+                await self.graph.set_node_status(node_id, NodeStatus.CANCELLED)
             finally:
                 # remove from running tasks in caller
                 pass

aethergraph/core/execution/global_scheduler.py

@@ -62,6 +62,7 @@ class RunState:
     ready_pending: set[str] = field(default_factory=set)  # nodes explicitly enqueued
     backoff_tasks: dict[str, asyncio.Task] = field(default_factory=dict)  # node_id -> sleeper task
     terminated: bool = False
+    cancelled: bool = False
 
     def capacity(self) -> int:
         return max(0, self.settings.max_concurrency - len(self.running_tasks))
@@ -163,6 +164,7 @@ class GlobalForwardScheduler:
         # mark runs as terminated & cancel sleepers/runners
         for rs in self._runs.values():
             rs.terminated = True
+            rs.cancelled = True
             for t in list(rs.backoff_tasks.values()):
                 t.cancel()
             for t in list(rs.running_tasks.values()):
@@ -185,6 +187,8 @@ class GlobalForwardScheduler:
         if not rs:
             return
         rs.terminated = True
+        rs.cancelled = True
+
         for t in list(rs.backoff_tasks.values()):
             t.cancel()
         for t in list(rs.running_tasks.values()):
@@ -240,20 +244,22 @@ class GlobalForwardScheduler:
                     and not tgt.resume_pending
                 ):
                     # compute a simple status
-                    status = "SUCCESS"
-                    for n in tgt.graph.nodes:
-                        if n.spec.type == "plan":
-                            continue
-                        if n.state.status == NodeStatus.FAILED:
-                            status = "FAILED"
-                            break
-                    await self._emit_run(
-                        RunEvent(
-                            run_id=tgt.run_id,
-                            status=status,
-                            timestamp=datetime.utcnow().timestamp(),
-                        )
+                    if tgt.cancelled:
+                        status = "CANCELLED"
+                    else:
+                        status = "SUCCESS"
+                        for n in tgt.graph.nodes:
+                            if n.spec.type == "plan":
+                                continue
+                            if n.state.status == NodeStatus.FAILED:
+                                status = "FAILED"
+                                break
+                    evt = RunEvent(
+                        run_id=tgt.run_id,
+                        status=status,
+                        timestamp=datetime.utcnow().timestamp(),
                     )
+                    await self._emit_run(evt)
                     break
 
             # 4) If nothing is running anywhere and nothing scheduled, decide how to wait
@@ -506,11 +512,11 @@ class GlobalForwardScheduler:
                     )
             except asyncio.CancelledError:
                 try:
-                    await rs.graph.set_node_status(node_id, NodeStatus.FAILED)
+                    await rs.graph.set_node_status(node_id, NodeStatus.CANCELLED)
                 except Exception as e:
                     if self._logger:
                         self._logger.warning(
-                            f"[GlobalForwardScheduler._start_node] failed to set node {node_id} as FAILED on cancellation: {e}"
+                            f"[GlobalForwardScheduler._start_node] failed to set node {node_id} as CANCELLED on cancellation: {e}"
                         )
             finally:
                 pass

aethergraph/core/execution/step_forward.py

@@ -3,6 +3,7 @@ import functools
 import inspect
 from typing import Any
 
+from aethergraph.contracts.services.channel import OutEvent
 from aethergraph.services.continuations.continuation import Continuation
 
 from ..graph.graph_refs import RESERVED_INJECTABLES  # {"context", "resume", "self"}
@@ -397,7 +398,15 @@ async def _enter_wait(
     # 3) Notify only if the tool hasn't already done it
     if not spec.get("notified", False):
         try:
-            await ctx.channels.notify(cont)
+            # TODO: This is a temporary fix. The proper way is to have the channel bus injected into the NodeContext
+            bus = node_ctx.services.channels
+            event = OutEvent(
+                type=cont.kind,
+                channel=cont.channel,
+                text=cont.prompt,
+                meta={"continuation_token": cont.token},
+            )
+            await bus.publish(event)
             if lg:
                 lg.debug("notified channel=%s", cont.channel)
         except Exception as e: