pydantic-graph 0.2.2__py3-none-any.whl → 1.24.0__py3-none-any.whl

pydantic_graph/_utils.py

@@ -1,18 +1,23 @@
 from __future__ import annotations as _annotations
 
 import asyncio
+import inspect
 import types
+import warnings
+from collections.abc import Callable, Generator
+from contextlib import contextmanager
 from functools import partial
-from typing import TYPE_CHECKING, Any, Callable, TypeVar
+from typing import TYPE_CHECKING, Any, TypeAlias, TypeVar, get_args, get_origin
 
-from logfire_api import LogfireSpan
-from typing_extensions import ParamSpec, TypeAlias, TypeIs, get_args, get_origin
+from logfire_api import Logfire, LogfireSpan
+from typing_extensions import ParamSpec, TypeIs
 from typing_inspection import typing_objects
 from typing_inspection.introspection import is_union_origin
 
 if TYPE_CHECKING:
     from opentelemetry.trace import Span
 
+_logfire = Logfire(otel_scope='pydantic-graph')
 
 AbstractSpan: TypeAlias = 'LogfireSpan | Span'
 
@@ -58,7 +63,7 @@ def get_union_args(tp: Any) -> tuple[Any, ...]:
     """Extract the arguments of a Union type if `response_type` is a union, otherwise return an empty tuple."""
     # similar to `pydantic_ai_slim/pydantic_ai/_result.py:get_union_args`
     if typing_objects.is_typealiastype(tp):
-        tp = tp.__value__
+        tp = tp.__value__  # pragma: no cover
 
     origin = get_origin(tp)
     if is_union_origin(origin):
@@ -96,8 +101,8 @@ def get_parent_namespace(frame: types.FrameType | None) -> dict[str, Any] | None
     If the graph is defined with generics `Graph[a, b]` then another frame is inserted, and we have to skip that
     to get the correct namespace.
     """
-    if frame is not None:
-        if back := frame.f_back:
+    if frame is not None:  # pragma: no branch
+        if back := frame.f_back:  # pragma: no branch
             if back.f_globals.get('__name__') == 'typing':
                 # If the class calling this function is generic, explicitly parameterizing the class
                 # results in a `typing._GenericAlias` instance, which proxies instantiation calls to the
@@ -135,3 +140,65 @@ async def run_in_executor(func: Callable[_P, _R], *args: _P.args, **kwargs: _P.k
         return await asyncio.get_running_loop().run_in_executor(None, partial(func, *args, **kwargs))
     else:
         return await asyncio.get_running_loop().run_in_executor(None, func, *args)  # type: ignore
+
+
+try:
+    from logfire._internal.config import (
+        LogfireNotConfiguredWarning,  # pyright: ignore[reportAssignmentType,reportPrivateImportUsage]
+    )
+except ImportError:  # pragma: lax no cover
+
+    class LogfireNotConfiguredWarning(UserWarning):
+        pass
+
+
+if TYPE_CHECKING:
+    logfire_span = _logfire.span
+else:
+
+    @contextmanager
+    def logfire_span(*args: Any, **kwargs: Any) -> Generator[LogfireSpan, None, None]:
+        """Create a Logfire span without warning if logfire is not configured."""
+        # TODO: Remove once Logfire has the ability to suppress this warning from non-user code
+        with warnings.catch_warnings():
+            warnings.filterwarnings('ignore', category=LogfireNotConfiguredWarning)
+            with _logfire.span(*args, **kwargs) as span:
+                yield span
+
+
+def infer_obj_name(obj: Any, *, depth: int) -> str | None:
+    """Infer the variable name of an object from the calling frame's scope.
+
+    This function examines the call stack to find what variable name was used
+    for the given object in the calling scope. This is useful for automatic
+    naming of objects based on their variable names.
+
+    Args:
+        obj: The object whose variable name to infer.
+        depth: Number of stack frames to traverse upward from the current frame.
+
+    Returns:
+        The inferred variable name if found, None otherwise.
+
+    Example:
+        Usage should generally look like `infer_name(self, depth=2)` or similar.
+    """
+    target_frame = inspect.currentframe()
+    if target_frame is None:
+        return None  # pragma: no cover
+    for _ in range(depth):
+        target_frame = target_frame.f_back
+        if target_frame is None:
+            return None
+
+    for name, item in target_frame.f_locals.items():
+        if item is obj:
+            return name
+
+    if target_frame.f_locals != target_frame.f_globals:  # pragma: no branch
+        # if we couldn't find the agent in locals and globals are a different dict, try globals
+        for name, item in target_frame.f_globals.items():
+            if item is obj:
+                return name
+
+    return None

pydantic_graph/beta/__init__.py

@@ -0,0 +1,25 @@
+"""The next version of the pydantic-graph framework with enhanced graph execution capabilities.
+
+This module provides a parallel control flow graph execution framework with support for:
+- 'Step' nodes for task execution
+- 'Decision' nodes for conditional branching
+- 'Fork' nodes for parallel execution coordination
+- 'Join' nodes and 'Reducer's for re-joining parallel executions
+- Mermaid diagram generation for graph visualization
+"""
+
+from .graph import Graph
+from .graph_builder import GraphBuilder
+from .node import EndNode, StartNode
+from .step import StepContext, StepNode
+from .util import TypeExpression
+
+__all__ = (
+    'EndNode',
+    'Graph',
+    'GraphBuilder',
+    'StartNode',
+    'StepContext',
+    'StepNode',
+    'TypeExpression',
+)

pydantic_graph/beta/decision.py

@@ -0,0 +1,276 @@
+"""Decision node implementation for conditional branching in graph execution.
+
+This module provides the Decision node type and related classes for implementing
+conditional branching logic in parallel control flow graphs. Decision nodes allow the graph
+to choose different execution paths based on runtime conditions.
+"""
+
+from __future__ import annotations
+
+import inspect
+from collections.abc import AsyncIterable, Callable, Iterable, Sequence
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Generic, get_origin
+
+from typing_extensions import Never, Self, TypeVar
+
+from pydantic_graph import BaseNode
+from pydantic_graph.beta.id_types import NodeID
+from pydantic_graph.beta.paths import Path, PathBuilder, TransformFunction
+from pydantic_graph.beta.step import NodeStep
+from pydantic_graph.beta.util import TypeOrTypeExpression
+from pydantic_graph.exceptions import GraphBuildingError
+
+if TYPE_CHECKING:
+    from pydantic_graph.beta.node_types import AnyDestinationNode, DestinationNode
+
+StateT = TypeVar('StateT', infer_variance=True)
+"""Type variable for graph state."""
+
+DepsT = TypeVar('DepsT', infer_variance=True)
+"""Type variable for graph dependencies."""
+
+HandledT = TypeVar('HandledT', infer_variance=True)
+"""Type variable used to track types handled by the branches of a Decision."""
+
+T = TypeVar('T', infer_variance=True)
+"""Generic type variable."""
+
+
+@dataclass(kw_only=True)
+class Decision(Generic[StateT, DepsT, HandledT]):
+    """Decision node for conditional branching in graph execution.
+
+    A Decision node evaluates conditions and routes execution to different
+    branches based on the input data type or custom matching logic.
+    """
+
+    id: NodeID
+    """Unique identifier for this decision node."""
+
+    branches: list[DecisionBranch[Any]]
+    """List of branches that can be taken from this decision."""
+
+    note: str | None
+    """Optional documentation note for this decision."""
+
+    def branch(self, branch: DecisionBranch[T]) -> Decision[StateT, DepsT, HandledT | T]:
+        """Add a new branch to this decision.
+
+        Args:
+            branch: The branch to add to this decision.
+
+        Returns:
+            A new Decision with the additional branch.
+        """
+        return Decision(id=self.id, branches=self.branches + [branch], note=self.note)
+
+    def _force_handled_contravariant(self, inputs: HandledT) -> Never:  # pragma: no cover
+        """Forces this type to be contravariant in the HandledT type variable.
+
+        This is an implementation detail of how we can type-check that all possible input types have
+        been exhaustively covered.
+
+        Args:
+            inputs: Input data of handled types.
+
+        Raises:
+            RuntimeError: Always, as this method should never be executed.
+        """
+        raise RuntimeError('This method should never be called, it is just defined for typing purposes.')
+
+
+SourceT = TypeVar('SourceT', infer_variance=True)
+"""Type variable for source data for a DecisionBranch."""
+
+
+@dataclass
+class DecisionBranch(Generic[SourceT]):
+    """Represents a single branch within a decision node.
+
+    Each branch defines the conditions under which it should be taken
+    and the path to follow when those conditions are met.
+
+    Note: with the current design, it is actually _critical_ that this class is invariant in SourceT for the sake
+    of type-checking that inputs to a Decision are actually handled. See the `# type: ignore` comment in
+    `tests.graph.beta.test_graph_edge_cases.test_decision_no_matching_branch` for an example of how this works.
+    """
+
+    source: TypeOrTypeExpression[SourceT]
+    """The expected type of data for this branch.
+
+    This is necessary for exhaustiveness-checking when handling the inputs to a decision node."""
+
+    matches: Callable[[Any], bool] | None
+    """An optional predicate function used to determine whether input data matches this branch.
+
+    If `None`, default logic is used which attempts to check the value for type-compatibility with the `source` type:
+    * If `source` is `Any` or `object`, the branch will always match
+    * If `source` is a `Literal` type, this branch will match if the value is one of the parametrizing literal values
+    * If `source` is any other type, the value will be checked for matching using `isinstance`
+
+    Inputs are tested against each branch of a decision node in order, and the path of the first matching branch is
+    used to handle the input value.
+    """
+
+    path: Path
+    """The execution path to follow when an input value matches this branch of a decision node.
+
+    This can include transforming, mapping, and broadcasting the output before sending to the next node or nodes.
+
+    The path can also include position-aware labels which are used when generating mermaid diagrams."""
+
+    destinations: list[AnyDestinationNode]
+    """The destination nodes that can be referenced by DestinationMarker in the path."""
+
+
+OutputT = TypeVar('OutputT', infer_variance=True)
+"""Type variable for the output data of a node."""
+
+NewOutputT = TypeVar('NewOutputT', infer_variance=True)
+"""Type variable for transformed output."""
+
+
+@dataclass(init=False)
+class DecisionBranchBuilder(Generic[StateT, DepsT, OutputT, SourceT, HandledT]):
+    """Builder for constructing decision branches with fluent API.
+
+    This builder provides methods to configure branches with destinations,
+    forks, and transformations in a type-safe manner.
+
+    Instances of this class should be created using [`GraphBuilder.match`][pydantic_graph.beta.graph_builder.GraphBuilder],
+    not created directly.
+    """
+
+    _decision: Decision[StateT, DepsT, HandledT]
+    """The parent decision node."""
+    _source: TypeOrTypeExpression[SourceT]
+    """The expected source type for this branch."""
+    _matches: Callable[[Any], bool] | None
+    """Optional matching predicate."""
+
+    _path_builder: PathBuilder[StateT, DepsT, OutputT]
+    """Builder for the execution path."""
+
+    def __init__(
+        self,
+        *,
+        decision: Decision[StateT, DepsT, HandledT],
+        source: TypeOrTypeExpression[SourceT],
+        matches: Callable[[Any], bool] | None,
+        path_builder: PathBuilder[StateT, DepsT, OutputT],
+    ):
+        # This manually-defined initializer is necessary due to https://github.com/python/mypy/issues/17623.
+        self._decision = decision
+        self._source = source
+        self._matches = matches
+        self._path_builder = path_builder
+
+    def to(
+        self,
+        destination: DestinationNode[StateT, DepsT, OutputT] | type[BaseNode[StateT, DepsT, Any]],
+        /,
+        *extra_destinations: DestinationNode[StateT, DepsT, OutputT] | type[BaseNode[StateT, DepsT, Any]],
+        fork_id: str | None = None,
+    ) -> DecisionBranch[SourceT]:
+        """Set the destination(s) for this branch.
+
+        Args:
+            destination: The primary destination node.
+            *extra_destinations: Additional destination nodes.
+            fork_id: Optional node ID to use for the resulting broadcast fork if multiple destinations are provided.
+
+        Returns:
+            A completed DecisionBranch with the specified destinations.
+        """
+        destination = get_origin(destination) or destination
+        extra_destinations = tuple(get_origin(d) or d for d in extra_destinations)
+        destinations = [(NodeStep(d) if inspect.isclass(d) else d) for d in (destination, *extra_destinations)]
+        return DecisionBranch(
+            source=self._source,
+            matches=self._matches,
+            path=self._path_builder.to(*destinations, fork_id=fork_id),
+            destinations=destinations,
+        )
+
+    def broadcast(
+        self, get_forks: Callable[[Self], Sequence[DecisionBranch[SourceT]]], /, *, fork_id: str | None = None
+    ) -> DecisionBranch[SourceT]:
+        """Broadcast this decision branch into multiple destinations.
+
+        Args:
+            get_forks: The callback that will return a sequence of decision branches to broadcast to.
+            fork_id: Optional node ID to use for the resulting broadcast fork.
+
+        Returns:
+            A completed DecisionBranch with the specified destinations.
+        """
+        fork_decision_branches = get_forks(self)
+        new_paths = [b.path for b in fork_decision_branches]
+        if not new_paths:
+            raise GraphBuildingError(f'The call to {get_forks} returned no branches, but must return at least one.')
+        path = self._path_builder.broadcast(new_paths, fork_id=fork_id)
+        destinations = [d for fdp in fork_decision_branches for d in fdp.destinations]
+        return DecisionBranch(source=self._source, matches=self._matches, path=path, destinations=destinations)
+
+    def transform(
+        self, func: TransformFunction[StateT, DepsT, OutputT, NewOutputT], /
+    ) -> DecisionBranchBuilder[StateT, DepsT, NewOutputT, SourceT, HandledT]:
+        """Apply a transformation to the branch's output.
+
+        Args:
+            func: Transformation function to apply.
+
+        Returns:
+            A new DecisionBranchBuilder where the provided transform is applied prior to generating the final output.
+        """
+        return DecisionBranchBuilder(
+            decision=self._decision,
+            source=self._source,
+            matches=self._matches,
+            path_builder=self._path_builder.transform(func),
+        )
+
+    def map(
+        self: DecisionBranchBuilder[StateT, DepsT, Iterable[T], SourceT, HandledT]
+        | DecisionBranchBuilder[StateT, DepsT, AsyncIterable[T], SourceT, HandledT],
+        *,
+        fork_id: str | None = None,
+        downstream_join_id: str | None = None,
+    ) -> DecisionBranchBuilder[StateT, DepsT, T, SourceT, HandledT]:
+        """Spread the branch's output.
+
+        To do this, the current output must be iterable, and any subsequent steps in the path being built for this
+        branch will be applied to each item of the current output in parallel.
+
+        Args:
+            fork_id: Optional ID for the fork, defaults to a generated value
+            downstream_join_id: Optional ID of a downstream join node which is involved when mapping empty iterables
+
+        Returns:
+            A new DecisionBranchBuilder where mapping is performed prior to generating the final output.
+        """
+        return DecisionBranchBuilder(
+            decision=self._decision,
+            source=self._source,
+            matches=self._matches,
+            path_builder=self._path_builder.map(fork_id=fork_id, downstream_join_id=downstream_join_id),
+        )
+
+    def label(self, label: str) -> DecisionBranchBuilder[StateT, DepsT, OutputT, SourceT, HandledT]:
+        """Apply a label to the branch at the current point in the path being built.
+
+        These labels are only used in generated mermaid diagrams.
+
+        Args:
+            label: The label to apply.
+
+        Returns:
+            A new DecisionBranchBuilder where the label has been applied at the end of the current path being built.
+        """
+        return DecisionBranchBuilder(
+            decision=self._decision,
+            source=self._source,
+            matches=self._matches,
+            path_builder=self._path_builder.label(label),
+        )