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

aethergraph/services/planning/orchestrator.py

@@ -0,0 +1,112 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from aethergraph.contracts.services.planning import (
+    IntentRouter,
+    SessionState,
+)
+from aethergraph.services.planning.plan_executor import (
+    ExecutionResult,
+    PlanExecutor,
+)
+from aethergraph.services.planning.plan_types import CandidatePlan
+from aethergraph.services.planning.planner import ActionPlanner
+
+from .planning_context_builder import PlanningContextBuilderProtocol
+from .quick_actions import QuickActionRegistry
+
+
+@dataclass
+class AgentTurnResult:
+    mode: str
+    message_to_user: str | None = None
+    plan: CandidatePlan | None = None
+    execution: ExecutionResult | None = None
+
+
+@dataclass
+class AgentOrchestrator:
+    router: IntentRouter
+    context_builder: PlanningContextBuilderProtocol
+    planner: ActionPlanner
+    executor: PlanExecutor
+    quick_actions: QuickActionRegistry
+
+    async def handle_turn(
+        self,
+        *,
+        user_message: str,
+        session_state: SessionState,
+    ) -> AgentTurnResult:
+        # 1) Route intent
+        routed = await self.router.route(
+            user_message=user_message,
+            session_state=session_state,
+        )
+
+        # 2) Mode dispatch
+        if routed.mode == "chat_only":
+            return AgentTurnResult(
+                mode="chat_only",
+                message_to_user="(chat-only mode: not yet implemented)",
+            )
+
+        if routed.mode == "quick_action":
+            handler = self.quick_actions.get_handler(routed.quick_action_id or "")
+            if handler is None:
+                return AgentTurnResult(
+                    mode="quick_action",
+                    message_to_user=f"Unknown quick action: {routed.quick_action_id}",
+                )
+            result = await handler(context={"user_message": user_message})
+            return AgentTurnResult(
+                mode="quick_action",
+                message_to_user=f"Quick action {routed.quick_action_id} done: {result!r}",
+            )
+
+        if routed.mode == "plan_and_execute":
+            # 3) Build planning context
+            planning_context = await self.context_builder.build(
+                user_message=user_message,
+                routed=routed,
+                session_state=session_state,
+            )
+
+            # 4) Plan
+            plan, history = await self.planner.plan_with_loop(planning_context)
+            if plan is None:
+                return AgentTurnResult(
+                    mode="plan_and_execute",
+                    message_to_user=(
+                        "I tried to build a plan but couldn't find a valid workflow. "
+                        "You may need to provide more details."
+                    ),
+                )
+
+            # 5) Execute
+            execution_result = await self.executor.execute(
+                plan,
+                user_inputs=planning_context.user_inputs,
+            )
+
+            if not execution_result.ok:
+                return AgentTurnResult(
+                    mode="plan_and_execute",
+                    message_to_user="The plan failed during execution.",
+                    plan=plan,
+                    execution=execution_result,
+                )
+
+            return AgentTurnResult(
+                mode="plan_and_execute",
+                message_to_user=f"Plan executed successfully. Outputs: {execution_result.outputs!r}",
+                plan=plan,
+                execution=execution_result,
+            )
+
+        # future unkonwn modes
+        return AgentTurnResult(
+            mode=routed.mode,
+            message_to_user=f"Unsupported mode: {routed.mode}",
+        )

aethergraph/services/planning/plan_executor.py

@@ -0,0 +1,506 @@
+# aethergraph/services/planning/plan_executor.py
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass, field
+import inspect
+from typing import TYPE_CHECKING, Any, Literal
+from uuid import uuid4
+
+from aethergraph.core.runtime.run_types import RunImportance, RunOrigin, RunVisibility
+from aethergraph.services.planning.action_catalog import ActionCatalog
+from aethergraph.services.planning.bindings import parse_binding
+
+if TYPE_CHECKING:
+    from aethergraph.api.v1.deps import RequestIdentity
+    from aethergraph.core.runtime.run_manager import RunManager
+    from aethergraph.services.planning.plan_types import CandidatePlan, ExecutionEventCallback
+
+ExecutionPhase = Literal[
+    "start",
+    "step_start",
+    "step_success",
+    "step_failure",
+    "success",
+    "failure",
+]
+
+
+@dataclass
+class ExecutionEvent:
+    """
+    Represents an event emitted during the execution of a plan.
+    This class is used to encapsulate information about the current state of
+    execution, which can be useful for logging purposes or updating the UI
+    with progress information.
+
+    Attributes:
+        phase (ExecutionPhase): The current phase of execution.
+        step_id (str | None): The identifier of the current step, if applicable.
+        message (str | None): An optional message providing additional context
+            about the event.
+        step_outputs (dict[str, Any] | None): Outputs produced by the step,
+            applicable for success phases.
+        error (Exception | None): The exception raised during execution,
+            applicable for failure phases.
+    """
+
+    phase: ExecutionPhase
+    step_id: str | None = None
+    message: str | None = None
+
+    # For success phases
+    step_outputs: dict[str, Any] | None = None
+
+    # For failures
+    error: Exception | None = None
+
+
+@dataclass
+class ExecutionResult:
+    """
+    Represents the result of an execution process.
+
+    Attributes:
+        ok (bool): Indicates whether the execution was successful.
+        outputs (dict[str, Any]): A dictionary containing the outputs of the execution.
+        errors (list[ExecutionEvent]): A list of errors that occurred during execution. Defaults to an empty list.
+    """
+
+    ok: bool
+    outputs: dict[str, Any]
+    errors: list[ExecutionEvent] = field(default_factory=list)
+
+
+@dataclass
+class BackgroundExecutionHandle:
+    """
+    BackgroundExecutionHandle is a data class that represents a handle for managing
+    the execution of a background task.
+
+    Attributes:
+        id (str): A unique identifier for the execution handle.
+        plan_id (str | None): The identifier of the associated plan, if any.
+        task (asyncio.Task[ExecutionResult]): The asyncio task representing the background execution.
+
+    Methods:
+        done() -> bool:
+            Checks if the background task has completed.
+        cancelled() -> bool:
+            Checks if the background task has been cancelled.
+        cancel() -> None:
+            Cancels the background task.
+        wait() -> ExecutionResult:
+            Awaits the completion of the background task and returns its result.
+    """
+
+    id: str
+    plan_id: str | None
+    task: asyncio.Task[ExecutionResult]
+
+    def done(self) -> bool:
+        return self.task.done()
+
+    def cancelled(self) -> bool:
+        return self.task.cancelled()
+
+    def cancel(self) -> None:
+        self.task.cancel()
+
+    async def wait(self) -> ExecutionResult:
+        return await self.task
+
+
+OnCompleteCallback = Callable[[ExecutionResult], Any | Awaitable[Any]]
+
+
+@dataclass
+class PlanExecutor:
+    """
+    Execute a validated CandidatePlan by invoking graphfns/graphs.
+
+    By default this uses run_async for each step (no RunStore / concurrency limits).
+    If a RunManager is provided, steps are executed via:
+
+        submit_run(...) + wait_run(..., return_outputs=True)
+
+    which:
+      - honors max_concurrent_runs
+      - creates RunRecords visible in the UI
+      - still exposes real Python outputs to the planner.
+
+    Assumptions:
+      - The plan has already been validated by FlowValidator.
+      - action_ref strings correspond to registry entries in the ActionCatalog,
+        using canonical refs like 'graph:foo@0.1.0' or 'graphfn:bar@0.1.0'.
+      - Bindings use the syntax:
+        * "${step_id.output_name}" for step outputs
+        * "${user.key}" for external/user inputs
+    """
+
+    catalog: ActionCatalog
+    # Optional: if provided, we use it instead of run_async for steps
+    run_manager: RunManager | None = None
+
+    # convenient access
+    @property
+    def registry(self):
+        return self.catalog.registry
+
+    async def execute(
+        self,
+        plan: CandidatePlan,
+        *,
+        user_inputs: dict[str, Any] | None = None,
+        on_event: ExecutionEventCallback | None = None,
+        # Optional execution context if using RunManager
+        identity: RequestIdentity | None = None,
+        visibility: RunVisibility | None = RunVisibility.normal,
+        importance: RunImportance | None = RunImportance.normal,
+        session_id: str | None = None,
+        agent_id: str | None = None,
+        app_id: str | None = None,
+        tags: list[str] | None = None,
+        origin: RunOrigin | None = None,
+    ) -> ExecutionResult:
+        """
+        Execute all steps in the plan in order.
+
+        Args:
+            plan: The candidate plan to execute (assumed validated).
+            user_inputs: Values that can be referenced as "${user.<key>}".
+            on_event: Optional callback to receive ExecutionEvent updates.
+            identity/session_id/agent_id/app_id/tags/origin:
+                Optional context passed through to RunManager when present.
+
+        Returns:
+            ExecutionResult:
+              - ok=True and final step outputs on success.
+              - ok=False and errors populated on first failure.
+        """
+        user_inputs = user_inputs or {}
+        step_results: dict[str, dict[str, Any]] = {}  # step_id -> outputs dict
+        errors: list[ExecutionEvent] = []
+        base_tags = list(tags or [])
+
+        await self._emit(
+            on_event,
+            ExecutionEvent(
+                phase="start",
+                message="Starting plan execution.",
+            ),
+        )
+
+        # Try to extract a plan id (best-effort, no hard dependency on field names)
+        plan_id = getattr(plan, "id", None) or getattr(plan, "plan_id", None)
+
+        # For now we assume steps are already in topological order (validator checked cycles).
+        for step in plan.steps:
+            await self._emit(
+                on_event,
+                ExecutionEvent(
+                    phase="step_start",
+                    step_id=step.id,
+                    message=f"Executing step '{step.id}' with action_ref='{step.action_ref}'.",
+                ),
+            )
+
+            try:
+                # Look up the underlying action object from the registry.
+                action_obj = self.registry.get(step.action_ref)
+
+                # Resolve inputs (bindings + literals)
+                bound_inputs = self._resolve_inputs(
+                    raw_inputs=step.inputs or {},
+                    step_results=step_results,
+                    user_inputs=user_inputs,
+                )
+
+                # Decide execution path:
+                # - If run_manager is None: use run_async (legacy/simple mode)
+                # - Else: use RunManager to honor concurrency + RunStore
+                if self.run_manager is None:
+                    from aethergraph.runner import run_async  # avoid top-level import cycle
+
+                    outputs = await run_async(action_obj, inputs=bound_inputs)
+                else:
+                    graph_id = self._graph_id_from_action_ref(step.action_ref)
+
+                    # Compose tags for this step-run
+                    step_tags = list(base_tags)
+                    if plan_id is not None:
+                        step_tags.append(f"plan:{plan_id}")
+                    step_tags.append(f"plan_step:{step.id}")
+
+                    # Spawn the run
+                    run_id = (
+                        f"plan-{plan_id or 'na'}-{session_id or 'na'}-{step.id}-{uuid4().hex[:8]}"
+                    )
+                    run_record = await self.run_manager.submit_run(
+                        graph_id=graph_id,
+                        inputs=bound_inputs,
+                        run_id=run_id,
+                        session_id=session_id,
+                        identity=identity,
+                        origin=origin,
+                        visibility=visibility,
+                        importance=importance,
+                        agent_id=agent_id,
+                        app_id=app_id,
+                        tags=step_tags,
+                    )
+                    # Wait for completion and grab real Python outputs
+                    finished_rec, outputs = await self.run_manager.wait_run(
+                        run_record.run_id,
+                        return_outputs=True,
+                    )
+
+                    # Interpret non-succeeded as failure
+                    status_str = getattr(finished_rec.status, "value", str(finished_rec.status))
+                    if status_str != "succeeded":
+                        # Mirror the run_async behaviour: raise so we hit the except below
+                        raise RuntimeError(
+                            f"Run for step '{step.id}' failed with status={status_str}, "
+                            f"error={finished_rec.error!r}"
+                        )
+
+                    # outputs may still be None if something went very wrong; guard it
+                    if outputs is None:
+                        outputs = {}
+
+                # Store step outputs for later bindings
+                step_results[step.id] = outputs
+
+                await self._emit(
+                    on_event,
+                    ExecutionEvent(
+                        phase="step_success",
+                        step_id=step.id,
+                        message=f"Step '{step.id}' completed.",
+                        step_outputs=outputs,
+                    ),
+                )
+
+            except Exception as exc:  # noqa: BLE001
+                failure_event = ExecutionEvent(
+                    phase="step_failure",
+                    step_id=step.id,
+                    message=f"Step '{step.id}' failed: {exc!r}",
+                    error=exc,
+                )
+                errors.append(failure_event)
+                await self._emit(on_event, failure_event)
+
+                # For v1: fail fast and mark the whole plan as failed.
+                final_failure = ExecutionEvent(
+                    phase="failure",
+                    step_id=step.id,
+                    message="Plan execution aborted due to step failure.",
+                    error=exc,
+                )
+                errors.append(final_failure)
+                await self._emit(on_event, final_failure)
+
+                return ExecutionResult(
+                    ok=False,
+                    outputs={},
+                    errors=errors,
+                )
+
+        # If we reach here, all steps succeeded.
+        final_step = plan.steps[-1]
+        final_outputs = step_results.get(final_step.id, {})
+
+        success_event = ExecutionEvent(
+            phase="success",
+            step_id=final_step.id,
+            message="Plan execution finished successfully.",
+            step_outputs=final_outputs,
+        )
+        await self._emit(on_event, success_event)
+
+        return ExecutionResult(
+            ok=True,
+            outputs=final_outputs,
+            errors=errors,
+        )
+
+    def execute_background(
+        self,
+        plan: CandidatePlan,
+        *,
+        user_inputs: dict[str, Any] | None = None,
+        on_event: ExecutionEventCallback | None = None,
+        identity: RequestIdentity | None = None,
+        visibility: RunVisibility | None = RunVisibility.normal,
+        importance: RunImportance | None = RunImportance.normal,
+        session_id: str | None = None,
+        agent_id: str | None = None,
+        app_id: str | None = None,
+        tags: list[str] | None = None,
+        origin: RunOrigin | None = None,
+        # callback when done
+        on_complete: Callable[[ExecutionResult], Any] | None = None,
+        # explicit id for tracing
+        exec_id: str | None = None,
+    ) -> BackgroundExecutionHandle:
+        """
+        Fire-and-forget convenience wrapper around execute().
+
+        - Schedules execute(...) on the current event loop using
+          asyncio.create_task()
+        - Returns immediately with a BackgroundExecutionHandle
+        - Still uses on_event for streaming progress
+        - Optionally triggers on_complete(result) when finished
+
+        NOTE: This does not change execution semantics; it's just
+        a lightweight scheduler helper. Orchestrator can use this
+        to let the user continue chatting while a plan runs.
+        """
+        loop = asyncio.get_event_loop()
+
+        plan_id = getattr(plan, "id", None) or getattr(plan, "plan_id", None)
+        exec_id = exec_id or f"bgexec-{plan_id or 'na'}-{uuid4().hex[:8]}"
+
+        async def _runner() -> ExecutionResult:
+            return await self.execute(
+                plan,
+                user_inputs=user_inputs,
+                on_event=on_event,
+                identity=identity,
+                visibility=visibility,
+                importance=importance,
+                session_id=session_id,
+                agent_id=agent_id,
+                app_id=app_id,
+                tags=tags,
+                origin=origin,
+            )
+
+        task: asyncio.Task[ExecutionResult] = loop.create_task(_runner(), name=exec_id)
+
+        if on_complete is not None:
+
+            def _done_callback(t: asyncio.Task[ExecutionResult]) -> None:
+                try:
+                    result = t.result()
+                except Exception as exc:
+                    # very defensive: crash in execute()
+                    # wrap the crash into ExecutionResult so caller get a structured error instead of an unhandled exception
+                    failure_evt = ExecutionEvent(
+                        phase="failure",
+                        step_id=None,
+                        message=f"Background execution '{exec_id}' failed: {t.exception()!r}",
+                        error=exc,
+                    )
+                    result = ExecutionResult(
+                        ok=False,
+                        outputs={},
+                        errors=[failure_evt],
+                    )
+                maybe_awaitable = on_complete(result)
+                if inspect.isawaitable(maybe_awaitable):
+                    # Fire and forget the completion callback as well
+                    loop.create_task(maybe_awaitable)
+
+            task.add_done_callback(_done_callback)
+
+        return BackgroundExecutionHandle(
+            id=exec_id,
+            plan_id=plan_id,
+            task=task,
+        )
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    async def _emit(
+        on_event: ExecutionEventCallback | None,
+        event: ExecutionEvent,
+    ) -> None:
+        if on_event is None:
+            return
+        try:
+            result = on_event(event)
+            if inspect.isawaitable(result):
+                await result
+        except Exception:
+            # Don't let logging/UI errors break execution
+            import logging
+
+            logger = logging.getLogger(__name__)
+            logger.warning("Error in execution on_event callback", exc_info=True)
+
+    @staticmethod
+    def _graph_id_from_action_ref(action_ref: str) -> str:
+        """
+        Extract the graph_id name from a canonical action_ref, e.g.:
+
+            "graph:foo@0.1.0"   -> "foo"
+            "graphfn:bar@0.1.0" -> "bar"
+
+        If the ref doesn't match this pattern, we fall back to the tail
+        after ":" (or the whole string).
+        """
+        ref = action_ref
+        if ":" in ref:
+            _, ref = ref.split(":", 1)
+        if "@" in ref:
+            ref, _ = ref.split("@", 1)
+        return ref
+
+    def _resolve_inputs(
+        self,
+        *,
+        raw_inputs: dict[str, Any],
+        step_results: dict[str, dict[str, Any]],
+        user_inputs: dict[str, Any],
+    ) -> dict[str, Any]:
+        """
+        Resolve all input values for a step:
+          - literals: kept as-is
+          - "${step_id.output_name}" -> previous step's outputs
+          - "${user.key}" -> user_inputs["key"]
+
+        We do this recursively, so bindings can appear inside nested dicts/lists.
+        """
+
+        def resolve_value(val: Any) -> Any:
+            # strings: may be bindings
+            if isinstance(val, str):
+                binding = parse_binding(val)
+
+                if binding.kind == "step_output":
+                    src_id = binding.source_step_id or ""
+                    out_name = binding.source_output_name or ""
+                    try:
+                        return step_results[src_id][out_name]
+                    except KeyError as exc:  # should be prevented by validator
+                        raise KeyError(
+                            f"Unknown step output reference: {src_id}.{out_name}"
+                        ) from exc
+
+                if binding.kind == "external":
+                    key = binding.external_key or ""
+                    if key not in user_inputs:
+                        raise KeyError(f"Missing user input for external binding: {key}")
+                    return user_inputs[key]
+
+                # literal or unsupported binding kinds: keep as-is
+                return val
+
+            # dict: resolve recursively
+            if isinstance(val, dict):
+                return {k: resolve_value(v) for k, v in val.items()}
+
+            # list/tuple: resolve recursively
+            if isinstance(val, (list, tuple)):  # noqa: UP038
+                return [resolve_value(v) for v in val]
+
+            # anything else: leave unchanged
+            return val
+
+        return {name: resolve_value(v) for name, v in raw_inputs.items()}