pydantic--pydantic-ai/tests/graph/builder/test_basenode_integration.py

"""Tests for using `BaseNode` subclasses inside the builder graph API."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Annotated, Any

import pytest

from pydantic_graph import BaseNode, End, GraphBuilder, GraphRunContext, StepContext, StepNode
from pydantic_graph.join import JoinNode, reduce_list_append

from ..._inline_snapshot import snapshot

pytestmark = pytest.mark.anyio


@dataclass
class IntegrationState:
    log: list[str] = field(default_factory=list[str])


async def test_basenode_in_builder_graph():
    """Test using BaseNode classes in a builder graph."""
    g = GraphBuilder(state_type=IntegrationState, input_type=int, output_type=str)

    @g.step
    async def prepare_input(ctx: StepContext[IntegrationState, object, int]) -> V1StartNode:
        ctx.state.log.append('V2Step: prepare')
        return V1StartNode(ctx.inputs + 1)

    @g.step
    async def process_result(ctx: StepContext[IntegrationState, object, str]) -> str:
        ctx.state.log.append('V2Step: process')
        return ctx.inputs.upper()

    @dataclass
    class V1StartNode(BaseNode[IntegrationState, object, str]):
        value: int

        async def run(self, ctx: GraphRunContext[IntegrationState, object]) -> V1MiddleNode:
            ctx.state.log.append(f'V1StartNode: {self.value}')
            return V1MiddleNode(self.value * 2)

    @dataclass
    class V1MiddleNode(BaseNode[IntegrationState, object, str]):
        value: int

        async def run(
            self, ctx: GraphRunContext[IntegrationState, object]
        ) -> Annotated[StepNode[IntegrationState, object], process_result]:
            ctx.state.log.append(f'V1MiddleNode: {self.value}')
            return process_result.as_node(f'Result: {self.value}')  # pyright: ignore[reportReturnType]  # TODO: GraphBuilder deps_type inference (v2 typevar default)

    g.add(
        g.node(V1StartNode),
        g.node(V1MiddleNode),
        g.edge_from(g.start_node).to(prepare_input),
        g.edge_from(process_result).to(g.end_node),
    )

    graph = g.build()
    state = IntegrationState()
    result = await graph.run(state=state, inputs=5)
    assert result == 'RESULT: 12'
    assert state.log == ['V2Step: prepare', 'V1StartNode: 6', 'V1MiddleNode: 12', 'V2Step: process']


async def test_step_to_basenode():
    """Test transitioning from a builder step to a BaseNode using StepNode."""
    g = GraphBuilder(state_type=IntegrationState, output_type=str)

    # `BaseNode` subclasses
    @dataclass
    class V1StartNode(BaseNode[IntegrationState, object, str]):
        value: int

        async def run(self, ctx: GraphRunContext[IntegrationState, object]) -> V1MiddleNode:  # pragma: no cover
            ctx.state.log.append(f'V1StartNode: {self.value}')
            return V1MiddleNode(self.value * 2)

    @dataclass
    class V1MiddleNode(BaseNode[IntegrationState, object, str]):
        value: int

        async def run(self, ctx: GraphRunContext[IntegrationState, object]) -> End[str]:  # pragma: no cover
            ctx.state.log.append(f'V1MiddleNode: {self.value}')
            return End(f'Result: {self.value}')

    @g.step
    async def step(
        ctx: StepContext[IntegrationState, object, None],
    ) -> V1StartNode:  # pragma: no cover
        ctx.state.log.append('V2Step')
        # Return a StepNode to transition to a BaseNode
        return V1StartNode(10)

    g.add(
        g.node(V1StartNode),
        g.node(V1MiddleNode),
        g.edge_from(g.start_node).to(step),
    )

    # Note: This will fail at type-checking but demonstrates the integration pattern
    # In practice, you'd need proper annotation handling


async def test_basenode_returning_basenode():
    """Test BaseNodes that return other BaseNodes."""

    @dataclass
    class FirstNode(BaseNode[IntegrationState, object, int]):
        value: int

        async def run(self, ctx: GraphRunContext[IntegrationState, object]) -> SecondNode:
            ctx.state.log.append('FirstNode')
            return SecondNode(self.value * 2)

    @dataclass
    class SecondNode(BaseNode[IntegrationState, object, int]):
        value: int

        async def run(self, ctx: GraphRunContext[IntegrationState, object]) -> End[int]:
            ctx.state.log.append('SecondNode')
            return End(self.value + 10)

    g = GraphBuilder(state_type=IntegrationState, input_type=int, output_type=int)

    @g.step
    async def create_first(ctx: StepContext[IntegrationState, object, int]) -> FirstNode:
        return FirstNode(ctx.inputs)

    g.add(
        g.node(FirstNode),
        g.node(SecondNode),
        g.edge_from(g.start_node).to(create_first),
    )

    graph = g.build()
    state = IntegrationState()
    result = await graph.run(state=state, inputs=5)
    assert result == 20  # 5 * 2 + 10
    assert state.log == ['FirstNode', 'SecondNode']


async def test_mixed_step_and_basenode_with_broadcast():
    """Test broadcasting with mixed step and `BaseNode` nodes."""
    g = GraphBuilder(state_type=IntegrationState, output_type=list[int])
    collect = g.join(reduce_list_append, initial_factory=list[int])

    @dataclass
    class ProcessNode(BaseNode[IntegrationState, object, Any]):
        value: int

        async def run(
            self, ctx: GraphRunContext[IntegrationState, object]
        ) -> Annotated[JoinNode[IntegrationState, object], collect]:
            ctx.state.log.append(f'ProcessNode: {self.value}')
            return collect.as_node(self.value * 2)  # pyright: ignore[reportReturnType]  # TODO: GraphBuilder deps_type inference (v2 typevar default)

    @g.step
    async def generate_values(ctx: StepContext[IntegrationState, object, None]) -> list[int]:
        return [1, 2, 3]

    @g.step
    async def create_node(ctx: StepContext[IntegrationState, object, int]) -> ProcessNode:
        return ProcessNode(ctx.inputs)

    g.add(
        g.node(ProcessNode),
        g.edge_from(g.start_node).to(generate_values),
        g.edge_from(generate_values).map().to(create_node),
        g.edge_from(collect).to(g.end_node),
    )

    graph = g.build()
    state = IntegrationState()
    result = await graph.run(state=state)
    assert sorted(result) == [2, 4, 6]
    assert len(state.log) == 3


async def test_basenode_type_hints_inferred():
    """Test that BaseNode type hints are properly inferred for edges."""

    @dataclass
    class StartNode(BaseNode[IntegrationState, object, str]):
        async def run(self, ctx: GraphRunContext[IntegrationState, object]) -> MiddleNode | End[str]:
            if ctx.state.log:
                return End('early exit')  # pragma: no cover
            ctx.state.log.append('StartNode')
            return MiddleNode()

    @dataclass
    class MiddleNode(BaseNode[IntegrationState, object, str]):
        async def run(self, ctx: GraphRunContext[IntegrationState, object]) -> End[str]:
            ctx.state.log.append('MiddleNode')
            return End('normal exit')

    g = GraphBuilder(state_type=IntegrationState, input_type=StartNode, output_type=str)

    g.add(
        g.node(StartNode),
        g.node(MiddleNode),
        g.edge_from(g.start_node).to(StartNode),
    )

    graph = g.build()
    state = IntegrationState()
    result = await graph.run(state=state, inputs=StartNode())
    assert result == 'normal exit'
    assert state.log == ['StartNode', 'MiddleNode']


async def test_basenode_conditional_return():
    """Test BaseNodes with conditional returns creating implicit decisions."""

    @dataclass
    class RouterNode(BaseNode[IntegrationState, object, str]):
        value: int

        async def run(self, ctx: GraphRunContext[IntegrationState, object]) -> PathA | PathB:
            if self.value < 10:
                return PathA()
            else:
                return PathB()

    @dataclass
    class PathA(BaseNode[IntegrationState, object, str]):
        async def run(self, ctx: GraphRunContext[IntegrationState, object]) -> End[str]:
            return End('Path A')

    @dataclass
    class PathB(BaseNode[IntegrationState, object, str]):
        async def run(self, ctx: GraphRunContext[IntegrationState, object]) -> End[str]:
            return End('Path B')

    g = GraphBuilder(state_type=IntegrationState, input_type=int, output_type=str)

    @g.step
    async def create_router(ctx: StepContext[IntegrationState, object, int]) -> RouterNode:
        return RouterNode(ctx.inputs)

    g.add(
        g.node(RouterNode),
        g.node(PathA),
        g.node(PathB),
        g.edge_from(g.start_node).to(create_router),
    )

    graph = g.build()

    assert str(graph) == snapshot("""\
stateDiagram-v2
  create_router
  RouterNode
  state decision <<choice>>
  PathA
  PathB

  [*] --> create_router
  create_router --> RouterNode
  RouterNode --> decision
  decision --> PathA
  decision --> PathB
  PathA --> [*]
  PathB --> [*]\
""")

    # Test path A
    result_a = await graph.run(state=IntegrationState(), inputs=5)
    assert result_a == 'Path A'

    # Test path B
    result_b = await graph.run(state=IntegrationState(), inputs=15)
    assert result_b == 'Path B'


async def test_match_node_with_base_node():
    """Test using match_node() to create decision branches for BaseNode classes

    Note: match_node is a complex API for integrating BaseNode types into decision logic.
    This test documents the intended usage pattern.
    """
    # This test is simplified to document match_node usage
    # The actual line coverage is achieved through internal graph construction
    pass