MessageNode to pydantic-graph Step Adapter Design¶

class MessageNode[TDeps, TResult](ABC):
    # Signals
    message_received = Signal[ChatMessage[Any]]()
    message_sent = Signal[ChatMessage[Any]]()

    # Lifecycle
    async def __aenter__(self) -> Self: ...
    async def __aexit__(self, ...): ...

    # Execution
    @abstractmethod
    async def run(self, *prompts: Any, **kwargs: Any) -> ChatMessage[TResult]: ...

    async def run_message(self, message: ChatMessage[Any], **kwargs: Any) -> ChatMessage[TResult]: ...

    @abstractmethod
    def run_iter(self, *prompts: Any, **kwargs: Any) -> AsyncIterator[ChatMessage[Any]]: ...

    async def run_stream(self, *prompts: PromptCompatible, ...) -> AsyncIterator[RichAgentStreamEvent[TResult]]: ...

    # Connections
    def connect_to(self, target: MessageNode | ProcessorCallback | Sequence[...], ...) -> Talk | TeamTalk: ...
    def __rshift__(self, other): ...  # syntactic sugar for connect_to

    # Context & State
    def get_context(self, data: Any = None, ...) -> NodeContext: ...
    @property
    def storage(self) -> StorageManager | None: ...

class BaseNode(ABC, Generic[StateT, DepsT, NodeRunEndT]):
    @abstractmethod
    async def run(self, ctx: GraphRunContext[StateT, DepsT]) -> BaseNode | End[NodeRunEndT]: ...

    @classmethod
    def get_node_id(cls) -> str: ...

@dataclass(init=False)
class StepContext(Generic[StateT, DepsT, InputT]):
    _state: StateT
    _deps: DepsT
    _inputs: InputT

    @property
    def state(self) -> StateT: ...
    @property
    def deps(self) -> DepsT: ...
    @property
    def inputs(self) -> InputT: ...

class StepFunction(Protocol[StateT, DepsT, InputT, OutputT]):
    def __call__(self, ctx: StepContext[StateT, DepsT, InputT]) -> Awaitable[OutputT]: ...

@dataclass(init=False)
class Step(Generic[StateT, DepsT, InputT, OutputT]):
    id: NodeID
    _call: StepFunction[StateT, DepsT, InputT, OutputT]
    label: str | None

    def as_node(self, inputs: InputT | None = None) -> StepNode: ...

@dataclass(init=False)
class Graph(Generic[StateT, DepsT, InputT, OutputT]):
    async def run(self, *, state: StateT = None, deps: DepsT = None, inputs: InputT = None) -> OutputT: ...

    @asynccontextmanager
    async def iter(self, *, state: StateT = None, deps: DepsT = None, inputs: InputT = None) -> AsyncIterator[GraphRun]: ...

class GraphRun(Generic[StateT, DepsT, OutputT]):
    async def next(self, value: EndMarker | Sequence[GraphTaskRequest] | None = None) -> EndMarker | Sequence[GraphTask]: ...
    @property
    def output(self) -> OutputT | None: ...

class MessageNode[TDeps, TResult](ABC):
    @property
    @abstractmethod
    def _step(self) -> Step[AgentPoolState, TDeps, ChatMessage[Any], ChatMessage[TResult]]: ...

    async def run(self, *prompts: Any, **kwargs: Any) -> ChatMessage[TResult]:
        graph = self._build_single_node_graph()
        state = AgentPoolState(node=self, prompts=prompts, kwargs=kwargs)
        result = await graph.run(state=state, deps=self._get_deps(), inputs=None)
        return result

    async def run_stream(self, *prompts: PromptCompatible, ...) -> AsyncIterator[RichAgentStreamEvent[TResult]]:
        graph = self._build_single_node_graph()
        state = AgentPoolState(node=self, prompts=prompts, kwargs=kwargs)
        async with graph.iter(state=state, deps=self._get_deps(), inputs=None) as graph_run:
            async for _ in graph_run:
                # Yield events captured from the node's internal event queue
                while event := self._event_queue.get_nowait():
                    yield event

class MessageNode[TDeps, TResult](BaseNode[AgentPoolState, TDeps, ChatMessage[TResult]], ABC):
    async def run(self, ctx: GraphRunContext[AgentPoolState, TDeps]) -> BaseNode | End[ChatMessage[TResult]]:
        # Execute the node, then return the next node (from connections) or End
        result = await self._execute(ctx.state.prompts)
        if next_node := self._resolve_next_node(result):
            return next_node
        return End(result)

class MessageNode[TDeps, TResult](ABC):
    @abstractmethod
    def build_step(self) -> Step[AgentPoolState, TDeps, ChatMessage[Any], ChatMessage[TResult]]: ...

    # run(), run_stream(), etc. are REMOVED from MessageNode
    # Execution happens only via Graph.run() / Graph.iter()

from dataclasses import dataclass, field
from typing import Any

from agentpool.messaging import ChatMessage


@dataclass
class AgentPoolState:
    """Shared state passed through the pydantic-graph execution.

    Holds the input prompts, node reference, and a conduit for
    streaming events back to the caller.
    """

    node: MessageNode[Any, Any]
    """The MessageNode being executed."""

    prompts: tuple[Any, ...]
    """Input prompts for this execution."""

    kwargs: dict[str, Any] = field(default_factory=dict)
    """Additional keyword arguments passed to run()."""

    # Streaming conduit
    event_queue: asyncio.Queue[RichAgentStreamEvent[Any]] = field(
        default_factory=asyncio.Queue
    )
    """Queue for streaming events from the Step back to run_stream()."""

    # Result capture
    result: ChatMessage[Any] | None = None
    """Final result populated by the Step upon completion."""

class MessageNode[TDeps, TResult](ABC):
    @property
    @abstractmethod
    def _step(self) -> Step[AgentPoolState, TDeps, ChatMessage[Any], ChatMessage[TResult]]:
        """Return the pydantic-graph Step representing this node's execution logic.

        The Step's call function receives a StepContext with:
        - state: AgentPoolState (prompts, event queue, node reference)
        - deps: TDeps (node dependencies, e.g., database connections)
        - inputs: ChatMessage[Any] (the input message, or None for root runs)

        The Step is responsible for:
        1. Emitting streaming events to state.event_queue
        2. Returning the final ChatMessage[TResult]
        3. Routing the result to connected nodes via state.node.connections

        Returns:
            A Step configured with this node's execution logic.
        """
        ...

class BaseAgent[TDeps = None, TResult = str](MessageNode[TDeps, TResult]):
    @property
    def _step(self) -> Step[AgentPoolState, TDeps, ChatMessage[Any], ChatMessage[TResult]]:
        async def _execute(ctx: StepContext[AgentPoolState, TDeps, ChatMessage[Any]]) -> ChatMessage[TResult]:
            state = ctx.state
            node = state.node
            assert isinstance(node, BaseAgent)

            # Reconstruct the run_stream() logic inside the Step
            # 1. Emit message_received signal
            user_msg = ChatMessage.user_prompt(message=state.prompts)
            await node.message_received.emit(user_msg)

            # 2. Stream events via the event queue
            async for event in node._stream_events_internal(
                prompts=state.prompts,
                user_msg=user_msg,
                deps=ctx.deps,
            ):
                await state.event_queue.put(event)

            # 3. Capture final result
            final_message = node._last_result  # set by _stream_events_internal
            state.result = final_message

            # 4. Emit message_sent signal
            await node.message_sent.emit(final_message)

            # 5. Route to connected nodes
            await node.connections.route_message(final_message)

            return final_message

        return Step(
            id=NodeID(self.name),
            call=_execute,
            label=self.description or self.name,
        )

class MessageNode[TDeps, TResult](ABC):
    def _build_single_node_graph(self) -> Graph[AgentPoolState, TDeps, ChatMessage[Any], ChatMessage[TResult]]:
        """Build a graph containing only this node's Step.

        The graph has:
        - start_node -> this node's Step -> end_node
        """
        builder = GraphBuilder[
            AgentPoolState, TDeps, ChatMessage[Any], ChatMessage[TResult]
        ](name=f"single_{self.name}")

        step = self._step
        builder.add_edge(builder.start_node, step)
        builder.add_edge(step, builder.end_node)

        return builder.build()

class MessageNode[TDeps, TResult](ABC):
    async def run(self, *prompts: Any, **kwargs: Any) -> ChatMessage[TResult]:
        """Execute node with prompts via pydantic-graph.

        Constructs a single-node graph and runs it to completion.
        """
        graph = self._build_single_node_graph()
        state = AgentPoolState(node=self, prompts=prompts, kwargs=kwargs)
        deps = self._get_deps()

        return await graph.run(state=state, deps=deps, inputs=None)

class MessageNode[TDeps, TResult](ABC):
    async def run_stream(
        self,
        *prompts: PromptCompatible,
        store_history: bool = True,
        message_id: str | None = None,
        session_id: str | None = None,
        parent_session_id: str | None = None,
        parent_id: str | None = None,
        message_history: MessageHistory | None = None,
        input_provider: InputProvider | None = None,
        wait_for_connections: bool | None = None,
        deps: TDeps | None = None,
        event_handlers: Sequence[AnyEventHandlerType] | None = None,
        depth: int = 0,
    ) -> AsyncIterator[RichAgentStreamEvent[TResult]]:
        """Run agent with streaming output via pydantic-graph Graph.iter().

        Uses Graph.iter() to drive execution step-by-step, yielding
        RichAgentStreamEvent tokens from the node's internal event queue.
        """
        graph = self._build_single_node_graph()
        state = AgentPoolState(
            node=self,
            prompts=prompts,
            kwargs={
                "store_history": store_history,
                "message_id": message_id,
                "session_id": session_id,
                "parent_session_id": parent_session_id,
                "parent_id": parent_id,
                "message_history": message_history,
                "input_provider": input_provider,
                "wait_for_connections": wait_for_connections,
                "deps": deps,
                "event_handlers": event_handlers,
                "depth": depth,
            },
        )

        async with graph.iter(state=state, deps=deps, inputs=None) as graph_run:
            async for _task in graph_run:
                # Drain the event queue after each graph step
                while not state.event_queue.empty():
                    try:
                        event = state.event_queue.get_nowait()
                        yield event
                    except asyncio.QueueEmpty:
                        break

        # Yield any remaining events after graph completion
        while not state.event_queue.empty():
            try:
                event = state.event_queue.get_nowait()
                yield event
            except asyncio.QueueEmpty:
                break

# message_received -> Step start
async def _execute(ctx: StepContext[AgentPoolState, TDeps, ChatMessage[Any]]) -> ChatMessage[TResult]:
    state = ctx.state
    node = state.node

    # Reconstruct the input message from prompts
    user_msg = ChatMessage.user_prompt(message=state.prompts)

    # Emit the signal (preserved for backward compatibility)
    await node.message_received.emit(user_msg)

    # ... execute ...

# message_sent -> Step completion (before returning)
    final_message = await _execute_core(...)
    await node.message_sent.emit(final_message)
    return final_message

# run_failed -> ErrorMarker recovery
# In run_stream(), when Graph.iter() yields an ErrorMarker:
async for marker in graph_run:
    if isinstance(marker, ErrorMarker):
        failed_event = BaseAgent.RunFailedEvent(
            agent_name=self.name,
            message="Agent stream failed",
            exception=marker.error,
        )
        await self.run_failed.emit(failed_event)
        # Allow pydantic-graph's error recovery or re-raise
        raise marker.error

# run_stream() mapped to Graph.iter() — detailed flow
async def run_stream(self, *prompts, **kwargs) -> AsyncIterator[RichAgentStreamEvent[TResult]]:
    graph = self._build_single_node_graph()
    state = AgentPoolState(node=self, prompts=prompts, kwargs=kwargs)

    async with graph.iter(state=state, deps=kwargs.get("deps"), inputs=None) as graph_run:
        # Drive the graph step by step
        while True:
            try:
                marker = await graph_run.next()
            except StopAsyncIteration:
                break

            if isinstance(marker, ErrorMarker):
                # Map to run_failed signal, then re-raise
                await self.run_failed.emit(...)
                raise marker.error

            # After each step, drain the event queue
            while not state.event_queue.empty():
                event = await asyncio.wait_for(state.event_queue.get(), timeout=0.01)
                yield event

    # Drain any final events
    while not state.event_queue.empty():
        yield state.event_queue.get_nowait()

class AgentPool:
    def _compile_connections_to_graph(self) -> Graph[AgentPoolState, Any, Any, Any] | None:
        """Compile all MessageNode connections into a pydantic-graph workflow.

        This is an optional optimization: if connections form a DAG,
        they can be executed via pydantic-graph's fork/join engine
        instead of AgentPool's ConnectionManager.
        """
        from pydantic_graph.graph_builder import GraphBuilder

        builder = GraphBuilder[AgentPoolState, Any, Any, Any](name="agent_pool_workflow")

        # Add all agent steps as nodes
        for name, node in self._registry.items():
            builder.add(node._step)

        # Add edges from connections
        for name, node in self._registry.items():
            for talk in node.connections.get_connections():
                for target in talk.targets:
                    builder.add_edge(
                        node._step,
                        target._step,
                        label=talk.connection_type,
                    )

        return builder.build()

class MessageNode:
    async def run(self, *prompts, **kwargs) -> ChatMessage[TResult]:
        result = await self._run_via_graph(prompts, kwargs)

        # After graph execution, route to connected nodes
        # This preserves the existing ConnectionManager behavior
        await self.connections.route_message(result)
        return result

# Team parallel -> Fork + Join
team = agent1 & agent2 & agent3
# Compiled to:
# fork -> [agent1, agent2, agent3] -> join -> result

# Team sequential -> Edge chain
team = agent1 | agent2 | agent3
# Compiled to:
# agent1 -> agent2 -> agent3 -> end

from __future__ import annotations

import asyncio
from collections.abc import AsyncIterator, Sequence
from dataclasses import dataclass, field
from typing import Any, Generic

from pydantic_graph.graph_builder import Graph, GraphBuilder
from pydantic_graph.step import Step, StepContext
from typing_extensions import TypeVar

from agentpool.agents.events import RichAgentStreamEvent
from agentpool.messaging import ChatMessage, MessageNode

StateT = TypeVar("StateT")
DepsT = TypeVar("DepsT")
InputT = TypeVar("InputT")
OutputT = TypeVar("OutputT")


@dataclass
class AgentPoolState:
    """State object passed through pydantic-graph execution."""

    node: MessageNode[Any, Any]
    prompts: tuple[Any, ...]
    kwargs: dict[str, Any] = field(default_factory=dict)
    event_queue: asyncio.Queue[RichAgentStreamEvent[Any]] = field(
        default_factory=asyncio.Queue
    )
    result: ChatMessage[Any] | None = None


class MessageNode[TDeps, TResult]:
    """Base class for all message processing nodes (adapted for pydantic-graph)."""

    @property
    @abstractmethod
    def _step(self) -> Step[AgentPoolState, TDeps, ChatMessage[Any], ChatMessage[TResult]]:
        """Return the pydantic-graph Step for this node."""
        ...

    def _build_single_node_graph(self) -> Graph[AgentPoolState, TDeps, ChatMessage[Any], ChatMessage[TResult]]:
        """Build a single-node graph wrapping this node's Step."""
        ...

    async def run(self, *prompts: Any, **kwargs: Any) -> ChatMessage[TResult]:
        """Execute via pydantic-graph Graph.run()."""
        ...

    async def run_stream(
        self, *prompts: Any, **kwargs: Any
    ) -> AsyncIterator[RichAgentStreamEvent[TResult]]:
        """Execute via pydantic-graph Graph.iter() with event streaming."""
        ...

    async def run_message(
        self, message: ChatMessage[Any], **kwargs: Any
    ) -> ChatMessage[TResult]:
        """Execute with a ChatMessage input."""
        ...