aethergraph 0.1.0a1__py3-none-any.whl

aethergraph/contracts/services/memory.py

@@ -0,0 +1,108 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Literal, Protocol, TypedDict
+
+EventKind = Literal[
+    "user_msg",
+    "assistant_msg",
+    "tool_start",
+    "tool_result",
+    "error",
+    "checkpoint",
+    "run_summary",
+    "rolling_summary",
+]
+
+
+@dataclass
+class Event:
+    event_id: str
+    ts: str
+    run_id: str
+    graph_id: str | None = None
+    node_id: str | None = None
+    agent_id: str | None = None
+    tool: str | None = None  # now used for tool topic: TODO: rename to topic in future
+    kind: EventKind = "tool_result"
+    stage: str | None = None
+    severity: int = 2
+    signal: float = 0.0
+    tags: list[str] | None = None
+    entities: list[str] | None = None
+    inputs: list[Value] | None = None
+    outputs: list[Value] | None = None
+    inputs_ref: dict[str, Any] | None = None
+    outputs_ref: dict[str, Any] | None = None
+    metrics: dict[str, float] | None = None
+    text: str | None = None
+    embedding: list[float] | None = None
+    pii_flags: dict[str, bool] | None = None
+    sources: list[str] | None = None
+    version: int = 1  # for schema evolution
+
+
+class HotLog(Protocol):
+    async def append(self, run_id: str, evt: Event, *, ttl_s: int, limit: int) -> None: ...
+    async def recent(
+        self, run_id: str, *, kinds: list[str] | None = None, limit: int = 50
+    ) -> list[Event]: ...
+
+
+class Persistence(Protocol):
+    async def append_event(self, run_id: str, evt: Event) -> None: ...
+    async def save_json(self, uri: str, obj: dict[str, Any]) -> None: ...
+
+
+class Indices(Protocol):
+    async def update(self, run_id: str, evt: Event) -> None: ...
+    async def last_by_name(self, run_id: str, name: str) -> dict[str, Any] | None: ...
+    async def latest_refs_by_kind(
+        self, run_id: str, kind: str, *, limit: int = 50
+    ) -> list[dict[str, Any]]: ...
+    async def last_outputs_by_topic(self, run_id: str, topic: str) -> dict[str, Any] | None: ...
+
+
+class Distiller(Protocol):
+    async def distill(
+        self, run_id: str, *, hotlog: HotLog, persistence: Persistence, indices: Indices, **kw
+    ) -> dict[str, Any]: ...
+
+
+# ---------- Vector Index and Embeddings Client Protocols ----------
+class VectorIndex(Protocol):
+    async def upsert(self, *, id: str, vector: list[float], metadata: dict) -> None: ...
+    async def delete(self, *, id: str) -> None: ...
+    async def query(
+        self, *, vector: list[float], k: int = 8, filter: dict | None = None
+    ) -> list[dict]: ...
+    async def flush(self) -> None: ...
+
+
+class EmbeddingsClient(Protocol):
+    async def embed_text(self, text: str, *, model: str | None = None) -> list[float]: ...
+    async def embed_texts(
+        self, texts: list[str], *, model: str | None = None
+    ) -> list[list[float]]: ...
+
+
+# ---------- I/O Value and Ref schemas ----------
+class Ref(TypedDict, total=False):
+    """A resolvable refernece to an external artifact or data."""
+
+    kind: str  # e.g. "spec", "design", "output", "tool_result"
+    uri: str  # e.g. "file://...", "mem://...", "db://..."
+    title: str | None  # optional human-readable title
+    mime: str | None  # optional MIME type, e.g. "image/png"
+
+
+class Value(TypedDict, total=False):
+    """
+    A named I/O slot that can hold any JSON-serializable value, including a Ref.
+    vtype declares the JSON type; if vtype == "ref", value must be a Ref dict.
+    """
+
+    name: str
+    vtype: Literal["ref", "number", "string", "boolean", "object", "array", "null"]
+    value: Any  # actual value; type depends on vtype
+    meta: dict[str, Any] | None  # optional metadata dictionary

aethergraph/contracts/services/resume.py

@@ -0,0 +1,28 @@
+from abc import abstractmethod
+from dataclasses import dataclass
+from typing import Protocol
+
+
+class ResumeBus(Protocol):
+    """
+    Abstract transport for resuming a waiting node.
+    Implementations may be:
+      - InProcessResumeBus: directly calls the in-memory scheduler
+      - HttpResumeBus: POSTs to a remote scheduler service (not shown here)
+    """
+
+    @abstractmethod
+    async def enqueue_resume(self, *, run_id: str, node_id: str, token: str, payload: dict) -> None:
+        """
+        Verify the continuation/token (or let the backend do it), then
+        trigger a resume for (run_id, node_id) with the given payload.
+        Should be idempotent and safe to call multiple times.
+        """
+        raise NotImplementedError
+
+
+@dataclass
+class ResumeEvent:
+    run_id: str
+    node_id: str
+    payload: dict

aethergraph/contracts/services/state_stores.py

@@ -0,0 +1,33 @@
+# aethergraph/persist/interfaces.py
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Protocol
+
+
+@dataclass
+class GraphSnapshot:
+    run_id: str
+    graph_id: str
+    rev: int
+    created_at: float  # epoch seconds
+    spec_hash: str  # detect spec drift
+    state: dict[str, Any]  # JSON-serializable TaskGraphState
+
+
+@dataclass
+class StateEvent:
+    run_id: str
+    graph_id: str
+    rev: int
+    ts: float
+    kind: str  # "STATUS" | "OUTPUT" | "INPUTS_BOUND" | "PATCH"
+    payload: dict[str, Any]
+
+
+class GraphStateStore(Protocol):
+    async def save_snapshot(self, snap: GraphSnapshot) -> None: ...
+    async def load_latest_snapshot(self, run_id: str) -> GraphSnapshot | None: ...
+    async def append_event(self, ev: StateEvent) -> None: ...
+    async def load_events_since(self, run_id: str, from_rev: int) -> list[StateEvent]: ...
+    async def list_run_ids(self, graph_id: str | None = None) -> list[str]: ...

aethergraph/contracts/services/wakeup.py

@@ -0,0 +1,28 @@
+# core/contracts/wakeup.py
+from dataclasses import dataclass
+from typing import Any, Protocol
+
+
+class WakeLease(Protocol):
+    """Lease handle for a wakeup message."""
+
+    id: str
+    msg: dict[str, Any]  # {run_id, node_id, token, payload, ...}
+    visibility_deadline: float
+
+
+class WakeupQueue(Protocol):
+    """Protocol for a wakeup queue service."""
+
+    async def enqueue(self, topic: str, msg: dict[str, Any], delay_s: float = 0) -> str: ...
+    async def lease(self, topic: str, max_items: int = 1, lease_s: int = 60) -> list[WakeLease]: ...
+    async def extend(self, lease: WakeLease, lease_s: int) -> None: ...
+    async def ack(self, lease: WakeLease) -> None: ...
+    async def nack(
+        self, lease: WakeLease, requeue_delay_s: float = 5
+    ) -> None: ...  # re-enqueue the message
+
+
+@dataclass
+class WakeupEvent:
+    node_id: str

aethergraph/core/execution/base_scheduler.py

@@ -0,0 +1,77 @@
+import asyncio
+from typing import Literal
+
+from aethergraph.core.graph.task_node import TaskNodeRuntime
+
+# from aethergraph.logging_config import logger
+
+
+ExecutionMode = Literal["forward", "backward"]
+
+
+class BaseScheduler:
+    def __init__(self, graph, mode: ExecutionMode):
+        self.graph = graph
+        self.mode = mode
+        self._pause_event = asyncio.Event()
+        self._pause_event.set()
+        self._terminated = False
+        self.running_tasks = {}  # Dictionary to track currently running tasks
+
+        self._nodes_to_pause = []  # Nodes that are requested to be paused
+
+    @property
+    def status(self) -> str:
+        if self._terminated:
+            return "terminated"
+        if not self._pause_event.is_set():
+            return "paused"
+        if self.get_running_task_node_ids():
+            return "running"
+        return "idle"
+
+    def reset_status(self):
+        """
+        Reset the scheduler's status to idle.
+        """
+        self._terminated = False
+        self._pause_event.set()
+        self.running_tasks.clear()
+        # logger.info(f"🔄 Scheduler status reset for graph `{self.graph.id}`")
+
+    def set_mode(self, mode: ExecutionMode):
+        """
+        Set the execution mode for this scheduler.
+        """
+        if mode not in ["forward", "backward"]:
+            raise ValueError(f"❌ Invalid execution mode: {mode}")
+        self.mode = mode
+        # logger.info(f"🔄 Execution mode set to {self.mode}")
+
+    def get_running_task_node_ids(self) -> list[str]:
+        """
+        Get a list of currently running task node IDs.
+        """
+        return [nid for nid, task in self.running_tasks.items() if not task.done()]
+
+    async def run(self):
+        raise NotImplementedError
+
+    async def run_from(self, node_ids: list[str]):
+        raise NotImplementedError("run_from() must be implemented in subclass")
+
+    async def pause(self):
+        self._pause_event.clear()
+
+    async def resume(self):
+        self._pause_event.set()
+
+    async def terminate(self):
+        self._terminated = True
+        raise NotImplementedError("terminate() must be implemented in subclass")
+
+    async def run_node(self, node: TaskNodeRuntime):
+        raise NotImplementedError("run_node() must be implemented in subclass")
+
+    async def step_next(self):
+        raise NotImplementedError("step_next() must be implemented in subclass")