aethergraph 0.1.0a1__py3-none-any.whl

aethergraph/core/graph/graph_fn.py

@@ -0,0 +1,219 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+import inspect
+
+from aethergraph.core.runtime.run_registration import RunRegistrationGuard
+
+from ..execution.retry_policy import RetryPolicy
+from ..runtime.runtime_env import RuntimeEnv
+from ..runtime.runtime_registry import current_registry  # ContextVar accessor
+from .graph_builder import graph  # context manager
+from .graph_refs import GRAPH_INPUTS_NODE_ID
+from .interpreter import Interpreter
+from .node_spec import TaskNodeSpec
+
+
+class GraphFunction:
+    def __init__(
+        self,
+        name: str,
+        fn: Callable,
+        inputs: list[str] | None = None,
+        outputs: list[str] | None = None,
+        version: str = "0.1.0",
+    ):
+        self.graph_id = name
+        self.name = name
+        self.fn = fn
+        self.inputs = inputs or []
+        self.outputs = outputs or []
+        self.version = version
+        self.registry_key: str | None = None
+        self.last_graph = None
+        self.last_context = None
+        self.last_memory_snapshot = None
+
+    async def run(
+        self,
+        *,
+        env: RuntimeEnv | None = None,
+        retry: RetryPolicy | None = None,
+        max_concurrency: int | None = None,
+        **inputs,
+    ):
+        """
+        Build a fresh TaskGraph and execute this function via the Interpreter.
+        If 'context' is declared in the fn signature, inject a NodeContext.
+        """
+        # Build env if not provided (use runner’s builder for consistency)
+        if env is None:
+            from ..runtime.graph_runner import _build_env  # internal helper
+
+            env, retry, max_concurrency = await _build_env(self, inputs)
+        if retry is None:
+            retry = RetryPolicy()
+
+        node_spec = TaskNodeSpec(
+            node_id=GRAPH_INPUTS_NODE_ID, type="inputs", metadata={"synthetic": True}
+        )
+        runtime_ctx = env.make_ctx(
+            node=node_spec, resume_payload=getattr(env, "resume_payload", None)
+        )
+        node_ctx = runtime_ctx.create_node_context(node=node_spec)
+
+        with graph(name=self.graph_id) as G:
+            interp = Interpreter(G, env, retry=retry, max_concurrency=max_concurrency)
+            run_id = env.run_id
+
+            # Register the scheduler for this run_id
+            with RunRegistrationGuard(
+                run_id=run_id, scheduler=interp.scheduler, container=env.container
+            ):
+                sig = inspect.signature(self.fn)
+                call_kwargs = dict(inputs)
+                if "context" in sig.parameters:
+                    call_kwargs["context"] = node_ctx
+
+                with interp.enter():
+                    res = self.fn(**call_kwargs)
+                    if inspect.isawaitable(res):
+                        res = await res
+
+            self.last_graph = G
+
+        res = _normalize_and_expose(G, res, self.outputs)
+        return res
+
+    # --- Syntactic sugar ---
+    async def __call__(self, **inputs):
+        """Async call to run the graph function.
+        Usage:
+           result = await my_graph_fn(input1=value1, input2=value2)
+        """
+        from ..runtime.graph_runner import run_async
+
+        return await run_async(self, inputs)
+
+    def sync(self, **inputs):
+        """Synchronous wrapper around async run(). Useful for quick tests or scripts.
+        Usage:
+           result = my_graph_fn.sync(input1=value1, input2=value2)
+        """
+        from ..runtime.graph_runner import run
+
+        return run(self, inputs)
+
+
+def _is_ref(x: object) -> bool:
+    return isinstance(x, dict) and x.get("_type") == "ref" and "from" in x and "key" in x
+
+
+def _is_nodehandle(x: object) -> bool:
+    return hasattr(x, "node_id") and hasattr(x, "output_keys")
+
+
+def _expose_from_handle(G, prefix: str, handle) -> dict:
+    oks = list(getattr(handle, "output_keys", []))
+    if not oks:
+        raise ValueError(f"NodeHandle '{getattr(handle, 'node_id', '?')}' has no output_keys")
+    out = {}
+    if prefix and len(oks) == 1:
+        # collapse single output to the provided key
+        k = oks[0]
+        ref = getattr(handle, k)
+        G.expose(prefix, ref)
+        out[prefix] = ref
+    else:
+        # multi-output (or top-level handle)
+        for k in oks:
+            key = f"{prefix}.{k}" if prefix else k
+            ref = getattr(handle, k)
+            G.expose(key, ref)
+            out[key] = ref
+    return out
+
+
+def _normalize_and_expose(G, ret, declared_outputs: list[str] | None) -> dict:
+    """
+    Normalize user return into {key: Ref or literal}.
+    - Dict of NodeHandles/Refs/literals supported
+    - Single NodeHandle supported
+    - Single literal supported (needs 1 declared output)
+    Also exposes Refs on G as boundary outputs.
+
+    Examples:
+    - return {"result": ref(...), "summary": node_handle(...), "count": 42}
+    - return node_handle(...)
+    """
+    result = {}
+
+    if isinstance(ret, dict):
+        for k, v in ret.items():
+            if _is_ref(v):
+                G.expose(k, v)
+                result[k] = v
+            elif _is_nodehandle(v):
+                result.update(_expose_from_handle(G, k, v))
+            else:
+                # literal stays literal; no expose
+                result[k] = v
+
+    elif _is_nodehandle(ret):
+        result.update(_expose_from_handle(G, "", ret))
+
+    else:
+        # single literal/ref case
+        if declared_outputs and len(declared_outputs) == 1:
+            key = declared_outputs[0]
+            if _is_ref(ret):
+                G.expose(key, ret)
+            result[key] = ret
+        else:
+            raise ValueError(
+                "Returning a single literal but outputs are not declared or >1. "
+                "Declare exactly one output or return a dict."
+            )
+
+    # If outputs were declared, restrict to those keys (keep order)
+    if declared_outputs:
+        result = {k: result[k] for k in declared_outputs if k in result}
+
+        # Validate presence
+        missing = [k for k in declared_outputs if k not in result]
+        if missing:
+            raise ValueError(f"Missing declared outputs: {missing}")
+
+    return result
+
+
+def graph_fn(
+    name: str,
+    inputs: list[str] | None = None,
+    outputs: list[str] | None = None,
+    version: str = "0.1.0",
+    agent: str | None = None,  # if agent is set, register this graph fn as an agent with given name
+) -> Callable[[Callable], GraphFunction]:
+    """Decorator to define a graph function."""
+
+    def decorator(fn: Callable):
+        gf = GraphFunction(name=name, fn=fn, inputs=inputs, outputs=outputs, version=version)
+        # Register in registry if given
+        registry = current_registry()
+
+        if registry is not None:
+            registry.register(
+                nspace="graphfn",
+                name=name,
+                version=version,
+                obj=gf,  # we register GraphFunction directly without spec -- graph function is already a runtime object
+            )
+
+        if agent:
+            assert (
+                registry is not None
+            ), "No registry available to register agent, make sure to have a current_registry() set up."
+            registry.register(nspace="agent", name=agent, version=version, obj=gf)
+        return gf
+
+    return decorator

aethergraph/core/graph/graph_io.py

@@ -0,0 +1,67 @@
+# ParamSpec, IOSpec, IOBindings, validators
+
+from dataclasses import asdict, dataclass, field
+from typing import Any, Literal
+
+from .graph_refs import normalize_binding
+
+
+@dataclass
+class ParamSpec:
+    """ParamSpec defines a single parameter's specification."""
+
+    schema: dict[str, Any] = field(default_factory=dict)  # JSON schema or empty
+    default: Any = None  # default value or None
+    source: Literal["arg", "ctx", "memory", "env", "secret", "kv"] | None = (
+        None  # where to bind from
+    )
+    doc: str | None = None  # optional description or docstring
+
+
+@dataclass
+class IOSpec:
+    required: dict[str, "ParamSpec"] = field(default_factory=dict)
+    optional: dict[str, "ParamSpec"] = field(default_factory=dict)
+    outputs: dict[str, "ParamSpec"] = field(default_factory=dict)
+
+    # Existing field (keep for back-compat)
+    expose: list[str] = field(default_factory=list)
+
+    # NEW: canonical bindings for exposed outputs (name -> Ref|literal)
+    expose_bindings: dict[str, Any] = field(default_factory=dict)
+
+    # ---- Convenience API (non-breaking) ----
+    def set_expose(self, name: str, binding: Any) -> None:
+        """Canonical way to record a public output and its binding."""
+        if name not in self.expose:
+            self.expose.append(name)
+        self.expose_bindings[name] = normalize_binding(binding)
+
+    def get_expose_names(self) -> list[str]:
+        # Use dict keys if present; otherwise fall back to list
+        if self.expose_bindings:
+            # ensure order is stable: preserve original list order if possible
+            ordered = [n for n in self.expose if n in self.expose_bindings]
+            # include any names defined only in bindings (edge cases)
+            ordered += [n for n in self.expose_bindings if n not in ordered]
+            return ordered
+        return list(self.expose)
+
+    def get_expose_bindings(self) -> dict[str, Any]:
+        # If only a list exists (legacy), return empty; caller can use heuristics if desired
+        return dict(self.expose_bindings)
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+
+@dataclass
+class IOBindings:
+    """IO bindings are used to bind actual values to the inputs/outputs defined in IOSpec."""
+
+    inbound: dict[str, str] = field(
+        default_factory=dict
+    )  # name -> source (arg, ctx, memory, env, secret, kv)
+    outbound: dict[str, str] = field(
+        default_factory=dict
+    )  # name -> destination (ctx, memory, kv, output)

aethergraph/core/graph/graph_refs.py

@@ -0,0 +1,154 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any, TypedDict
+
+GRAPH_INPUTS_NODE_ID = "__graph_inputs__"  # special node_id for graph inputs
+RESERVED_INJECTABLES = {"resume", "context", "self"}
+
+
+REF_TYPE = "ref"
+
+
+class RefDict(TypedDict):
+    _type: str
+    from_: str  # 'from' is reserved in Python keyword args, keep key as 'from' in payload though
+    key: str
+
+
+Ref = dict[str, str]  # {"_type": "ref", "from": "<node_id>", "key": "<output_key>"}
+
+# ---------- Constructors ----------
+
+
+def ref(node_id: str, key: str) -> Ref:
+    return {"_type": "ref", "from": node_id, "key": key}
+
+
+def arg(name: str) -> Ref:
+    return ref(GRAPH_INPUTS_NODE_ID, name)
+
+
+# ---------- Type checks / Normalizations ----------
+def is_ref(x: Any) -> bool:
+    """True if x is a dict that looks like a Ref."""
+    return isinstance(x, Mapping) and x.get("_type") == REF_TYPE and "from" in x and "key" in x
+
+
+def is_arg_ref(x: Any) -> bool:
+    """True if x is a ref pointing to __graph_inputs__."""
+    return is_ref(x) and x.get("from") == GRAPH_INPUTS_NODE_ID
+
+
+def to_tuple(x: Ref | tuple[str, str]) -> tuple[str, str] | None:
+    """Return (node_id, key) if x is a Ref/tuple; else None."""
+    if isinstance(x, tuple) and len(x) == 2 and all(isinstance(s, str) for s in x):
+        return x  # already canonical enough
+    if is_ref(x):
+        return x["from"], x["key"]
+    return None
+
+
+def from_tuple(node_key: tuple[str, str]) -> Ref:
+    """Build a Ref from (node_id, key)."""
+    node_id, key = node_key
+    return ref(node_id, key)
+
+
+def normalize_binding(x: Any) -> Any:
+    """
+    Normalize a binding value (Ref | (node, key) | literal) into Ref|literal.
+    - If tuple -> Ref
+    - If Ref -> ensure minimal canonical shape
+    - Else literal passthrough
+    """
+    t = to_tuple(x)
+    if t is not None:
+        return from_tuple(t)
+    if is_ref(x):
+        # keep only the keys we care about (defensive)
+        return {"_type": REF_TYPE, "from": x["from"], "key": x["key"]}
+    return x  # literal
+
+
+# ---------- Resolution ----------
+def resolve_ref(reference: Ref, node_outputs: Mapping[str, Mapping[str, Any]]) -> Any:
+    """
+    Resolve a Ref against the current node_outputs: {node_id: {output_key: value}}.
+    Returns None if missing.
+    """
+    src = reference["from"]
+    key = reference["key"]
+
+    if src is None or key is None:
+        raise KeyError(f"Bad Ref: {ref}")
+    if src not in node_outputs:
+        raise KeyError(f"Upstream node '{src}' has no outputs yet")
+    if key not in node_outputs[src]:
+        raise KeyError(f"Output '{key}' not found on node '{src}'")
+
+    outs = node_outputs.get(src)
+    return outs.get(key) if isinstance(outs, Mapping) else None
+
+
+def resolve_any(val, *, graph_inputs: dict[str, Any], outputs_by_node: dict[str, dict[str, Any]]):
+    """Recursively resolve any value that may contain Refs or Args. This function is used
+    to resolve inputs for a node before execution.
+    Args:
+        val: The value to resolve. Can be a literal, dict, list, or Ref/Arg.
+        graph_inputs: The dict of graph inputs for Arg resolution.
+        outputs_by_node: The dict of node_id to outputs for Ref resolution.
+    Returns:
+        The fully resolved value.
+    """
+    # Arg shape: {"_type":"arg","key":"<input_key>"}
+    if isinstance(val, dict):
+        t = val.get("_type")
+        if t == "arg":
+            k = val.get("key")
+            if k not in graph_inputs:
+                raise KeyError(f"Graph input '{k}' was not provided")
+            return graph_inputs[k]
+        if t == "ref":
+            return resolve_ref(val, outputs_by_node)
+        # regular dict → recurse
+        return {
+            k: resolve_any(v, graph_inputs=graph_inputs, outputs_by_node=outputs_by_node)
+            for k, v in val.items()
+        }
+    if isinstance(val, list | tuple):
+        cast = list if isinstance(val, list) else tuple
+        return cast(
+            resolve_any(v, graph_inputs=graph_inputs, outputs_by_node=outputs_by_node) for v in val
+        )
+    return val  # literal
+
+
+def resolve_binding(binding: Any, node_outputs: Mapping[str, Mapping[str, Any]]) -> Any:
+    """
+    Resolve a binding that can be Ref or literal. Literals pass through unchanged.
+    """
+    if is_ref(binding):
+        return resolve_ref(binding, node_outputs)
+    return binding
+
+
+# ---------- Pretty helpers ----------
+
+
+def ref_str(x: Ref | tuple[str, str] | Any) -> str:
+    """Human-friendly string for logs."""
+    t = to_tuple(x)
+    if t is None:
+        return repr(x)
+    node_id, key = t
+    return f"{node_id}.{key}"
+
+
+# --------- Marker checks ----------
+def is_arg_marker(x: Any) -> bool:
+    return isinstance(x, Mapping) and x.get("_type") == "arg" and "key" in x
+
+
+def is_context_marker(x: Any) -> bool:
+    return isinstance(x, Mapping) and x.get("_type") == "context" and "key" in x

aethergraph/core/graph/graph_spec.py

@@ -0,0 +1,115 @@
+from dataclasses import dataclass, field
+import inspect
+from typing import Any
+
+from .graph_io import IOBindings, IOSpec
+from .node_spec import TaskNodeSpec
+
+
+@dataclass
+class TaskGraphSpec:
+    graph_id: str
+    version: str = "0.1.0"
+    nodes: dict[str, TaskNodeSpec] = field(default_factory=dict)  # node_id -> TaskNodeSpec
+    io: IOSpec = field(default_factory=IOSpec)  # inputs/outputs
+    bindings: IOBindings | None = None  # input/output bindings
+    meta: dict[str, Any] = field(default_factory=dict)  # additional metadata
+
+    def canonical(self) -> str:
+        return f"graph:{self.graph_id}@{self.version}"
+
+    @property
+    def inputs_required(self) -> set[str]:
+        return set(self.io.required.keys())
+
+    @property
+    def inputs_optional(self) -> dict[str, Any]:
+        return {k: p.default for k, p in self.io.optional.items()}
+
+    @property
+    def outputs(self) -> dict[str, Any]:
+        return {k: p.default for k, p in self.io.outputs.items()}
+
+    def io_summary_lines(self) -> list[str]:
+        return [
+            f"required: {_fmt_set(self.inputs_required)}",
+            f"optional: {_fmt_opt_map(self.inputs_optional, show_values=False)}",
+            f"outputs:  {_fmt_outputs_map(self.outputs)}",
+        ]
+
+
+@dataclass
+class GraphView:
+    """A read-only view of the graph's spec and state."""
+
+    graph_id: str
+    nodes: dict[str, Any]  # node_id -> TaskNodeRuntime
+    metadata: dict[str, Any] = field(default_factory=dict)  # Optional metadata
+
+    # helpers, no mutation
+    def get_dependents(self, nid: str) -> list[str]:
+        """Get list of node_ids that depend on the given node_id."""
+        return [x.node_id for x in self.nodes.values() if nid in x.dependencies]
+
+    def get_root_nodes(self) -> list[str]:
+        """Get list of root node_ids (no dependencies)."""
+        return [x.node_id for x in self.nodes.values() if not x.dependencies]
+
+
+# ---------- helpers for printing and debugging ----------
+
+
+def _short(x: Any, maxlen: int = 42) -> str:
+    s = str(x)
+    return s if len(s) <= maxlen else s[: maxlen - 1] + "…"
+
+
+def _status_label(s: Any) -> str:
+    # Accept Enum-like (with .name), strings, or None
+    if s is None:
+        return "-"
+    return getattr(s, "name", str(s))
+
+
+def _logic_label(logic: Any) -> str:
+    # Show a dotted path when possible; fall back to repr/str
+    if isinstance(logic, str):
+        return logic
+    # Unwrap @tool proxies if present
+    impl = getattr(logic, "__aether_impl__", logic)
+    if inspect.isfunction(impl) or inspect.ismethod(impl):
+        mod = getattr(impl, "__module__", None) or ""
+        name = getattr(impl, "__name__", None) or "tool"
+        return f"{mod}.{name}".strip(".")
+    return _short(repr(logic), 80)
+
+
+def _fmt_set(xs: set | None) -> str:
+    return ", ".join(sorted(map(str, xs))) if xs else "—"
+
+
+def _fmt_opt_map(d: dict | None, *, show_values: bool = False, maxval: int = 26) -> str:
+    if not d:
+        return "—"
+    if show_values:
+        items = [f"{k}={_short(v, maxval)}" for k, v in d.items()]
+    else:
+        items = list(map(str, d.keys()))
+    return ", ".join(sorted(items)) if items else "—"
+
+
+def _fmt_outputs_map(d: dict | None) -> str:
+    """
+    Show graph outputs mapping; if a value looks like a Ref(node_id, key),
+    render as 'out_key ← node_id.key'. Otherwise, just list keys.
+    """
+    if not d:
+        return "—"
+    parts = []
+    for out_k, v in d.items():
+        # duck-typed Ref
+        if hasattr(v, "node_id") and hasattr(v, "key"):
+            parts.append(f"{out_k} ← {v.node_id}.{v.key}")
+        else:
+            parts.append(str(out_k))
+    return ", ".join(sorted(parts))

aethergraph/core/graph/graph_state.py

@@ -0,0 +1,59 @@
+from dataclasses import dataclass, field
+from typing import Any
+
+from .node_spec import TaskNodeSpec
+from .node_state import TaskNodeState
+
+
+@dataclass
+class GraphPatch:
+    # Used for mutation of a graph's topology; NOT USED YET
+    op: str  # "add_node", "remove_node", "update_node", "add_edge", "remove_edge"
+    payload: dict[str, Any]  # details depend on op type -> To be defined later
+
+
+@dataclass
+class TaskGraphState:
+    run_id: str | None = (
+        None  # unique run identifier, used when in agent or program execution, and run_id is known from agent/program level
+    )
+    nodes: dict[str, TaskNodeState] = field(default_factory=dict)  # node_id -> TaskNodeState
+    # node_status: Dict[str, str] = field(default_factory=dict)  # node_id -> status ("pending", "running", "completed", "failed", etc.)
+    # node_outputs: Dict[str, Any] = field(default_factory=dict)  # node_id -> output data
+    _bound_inputs: dict[str, Any] | None = field(
+        default=None, repr=False
+    )  # inputs bound at runtime
+    rev: int = 0  # revision number, incremented on each mutation
+    patches: list[GraphPatch] = field(
+        default_factory=list, repr=False
+    )  # list of patches applied to the graph
+
+    def default_node_states(self, spec: TaskNodeSpec):
+        # Initialize node states based on the given spec
+        for nid in spec.nodes:
+            if nid not in self.nodes:
+                self.nodes[nid] = TaskNodeState()
+                self.node_status[nid] = "PENDING"
+
+    def summary_line(self) -> str:
+        from collections import Counter
+
+        sc = Counter(self.node_status.values())
+        counts = ", ".join(f"{k}={v}" for k, v in sorted(sc.items())) or "—"
+        bound = list(self._bound_inputs.keys()) if self._bound_inputs else "—"
+        return (
+            f"bound_inputs={bound}, node_outputs={len(self.node_outputs)}, status_counts: {counts}"
+        )
+
+    @property
+    def node_statuses(self) -> dict[str, str]:
+        return {nid: ns.status for nid, ns in self.nodes.items()}
+
+    # alias to node_statuses
+    @property
+    def node_status(self) -> dict[str, str]:
+        return {nid: ns.status for nid, ns in self.nodes.items()}
+
+    @property
+    def node_outputs(self) -> dict[str, Any]:
+        return {nid: ns.outputs for nid, ns in self.nodes.items() if ns.outputs}