diff --git a/docs/docs/reference/graphs.md b/docs/docs/reference/graphs.md index 3c25180ca..779f95095 100644 --- a/docs/docs/reference/graphs.md +++ b/docs/docs/reference/graphs.md @@ -69,8 +69,12 @@ builder.add_conditional_edges("my_node", my_condition) ## Send -::: langgraph.pregel.types.Send +::: langgraph.types.Send + +## Interrupt + +::: langgraph.types.Interrupt ## RetryPolicy -::: langgraph.pregel.types.RetryPolicy +::: langgraph.types.RetryPolicy diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index e17ce7de8..ef4a8a486 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -1,20 +1,9 @@ from types import MappingProxyType from typing import Any, Mapping +from langgraph.types import Interrupt, Send # noqa: F401 # Interrupt, Send re-exported for backwards compatibility -def __getattr__(name: str) -> Any: - if name in globals(): - return globals()[name] - elif name == "Interrupt": - from langgraph.pregel.types import Interrupt - - return Interrupt - elif name == "Send": - from langgraph.pregel.types import Send - - return Send - raise AttributeError(f"module {__name__} has no attribute {name}") # --- Empty read-only containers --- @@ -54,6 +43,8 @@ CONFIG_KEY_STORE = "__pregel_store" # holds a `BaseStore` made available to managed values CONFIG_KEY_RESUMING = "__pregel_resuming" # holds a boolean indicating if subgraphs should resume from a previous checkpoint +CONFIG_KEY_GRAPH_COUNT = "__pregel_graph_count" +# holds the number of subgraphs executed in a given task, used to raise errors CONFIG_KEY_TASK_ID = "__pregel_task_id" # holds the task ID for the current task CONFIG_KEY_DEDUPE_TASKS = "__pregel_dedupe_tasks" diff --git a/libs/langgraph/langgraph/errors.py b/libs/langgraph/langgraph/errors.py index 08bed5927..c7c5a518a 100644 --- a/libs/langgraph/langgraph/errors.py +++ b/libs/langgraph/langgraph/errors.py @@ -1,7 +1,7 @@ from typing import Any, Sequence from langgraph.checkpoint.base import EmptyChannelError # noqa: F401 -from langgraph.constants import Interrupt +from langgraph.types import Interrupt # EmptyChannelError re-exported for backwards compatibility diff --git a/libs/langgraph/langgraph/graph/graph.py b/libs/langgraph/langgraph/graph/graph.py index c5a043ee7..a5a4db3ac 100644 --- a/libs/langgraph/langgraph/graph/graph.py +++ b/libs/langgraph/langgraph/graph/graph.py @@ -38,8 +38,8 @@ from langgraph.constants import ( from langgraph.errors import InvalidUpdateError from langgraph.pregel import Channel, Pregel from langgraph.pregel.read import PregelNode -from langgraph.pregel.types import All from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry +from langgraph.types import All from langgraph.utils.runnable import RunnableCallable, coerce_to_runnable logger = logging.getLogger(__name__) diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index cee0fb849..4fdff2a49 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -45,9 +45,9 @@ from langgraph.managed.base import ( is_writable_managed_value, ) from langgraph.pregel.read import ChannelRead, PregelNode -from langgraph.pregel.types import All, RetryPolicy from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry from langgraph.store.base import BaseStore +from langgraph.types import All, RetryPolicy from langgraph.utils.fields import get_field_default from langgraph.utils.pydantic import create_model from langgraph.utils.runnable import coerce_to_runnable diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index f1a9aa578..c9fec2e57 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -83,11 +83,11 @@ from langgraph.pregel.messages import StreamMessagesHandler from langgraph.pregel.read import PregelNode from langgraph.pregel.retry import RetryPolicy from langgraph.pregel.runner import PregelRunner -from langgraph.pregel.types import All, StateSnapshot, StreamMode from langgraph.pregel.utils import get_new_channel_versions from langgraph.pregel.validate import validate_graph, validate_keys from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry from langgraph.store.base import BaseStore +from langgraph.types import All, StateSnapshot, StreamMode from langgraph.utils.config import ( ensure_config, merge_configs, diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 44da0bbd5..f9b0096c8 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -51,7 +51,7 @@ from langgraph.pregel.io import read_channel, read_channels from langgraph.pregel.log import logger from langgraph.pregel.manager import ChannelsManager from langgraph.pregel.read import PregelNode -from langgraph.pregel.types import All, PregelExecutableTask, PregelTask +from langgraph.types import All, PregelExecutableTask, PregelTask from langgraph.utils.config import merge_configs, patch_config GetNextVersion = Callable[[Optional[V], BaseChannel], V] diff --git a/libs/langgraph/langgraph/pregel/debug.py b/libs/langgraph/langgraph/pregel/debug.py index 9c4661c0c..982182842 100644 --- a/libs/langgraph/langgraph/pregel/debug.py +++ b/libs/langgraph/langgraph/pregel/debug.py @@ -22,7 +22,7 @@ from langgraph.channels.base import BaseChannel from langgraph.checkpoint.base import Checkpoint, CheckpointMetadata, PendingWrite from langgraph.constants import ERROR, INTERRUPT, TAG_HIDDEN from langgraph.pregel.io import read_channels -from langgraph.pregel.types import PregelExecutableTask, PregelTask, StateSnapshot +from langgraph.types import PregelExecutableTask, PregelTask, StateSnapshot class TaskPayload(TypedDict): diff --git a/libs/langgraph/langgraph/pregel/io.py b/libs/langgraph/langgraph/pregel/io.py index 6542b1d91..ef9822641 100644 --- a/libs/langgraph/langgraph/pregel/io.py +++ b/libs/langgraph/langgraph/pregel/io.py @@ -5,7 +5,7 @@ from langchain_core.runnables.utils import AddableDict from langgraph.channels.base import BaseChannel, EmptyChannelError from langgraph.constants import EMPTY_SEQ, ERROR, INTERRUPT, TAG_HIDDEN from langgraph.pregel.log import logger -from langgraph.pregel.types import PregelExecutableTask +from langgraph.types import PregelExecutableTask def read_channel( diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index c97cc8a69..7eec59575 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -94,10 +94,10 @@ from langgraph.pregel.io import ( ) from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager from langgraph.pregel.read import PregelNode -from langgraph.pregel.types import All, PregelExecutableTask, StreamMode from langgraph.pregel.utils import get_new_channel_versions from langgraph.store.base import BaseStore from langgraph.store.batch import AsyncBatchedStore +from langgraph.types import All, PregelExecutableTask, StreamMode from langgraph.utils.config import patch_configurable V = TypeVar("V") diff --git a/libs/langgraph/langgraph/pregel/retry.py b/libs/langgraph/langgraph/pregel/retry.py index 90ccaa7d0..476b8ef32 100644 --- a/libs/langgraph/langgraph/pregel/retry.py +++ b/libs/langgraph/langgraph/pregel/retry.py @@ -4,9 +4,9 @@ import random import time from typing import Optional, Sequence -from langgraph.constants import CONFIG_KEY_RESUMING +from langgraph.constants import CONFIG_KEY_GRAPH_COUNT, CONFIG_KEY_RESUMING from langgraph.errors import GraphInterrupt -from langgraph.pregel.types import PregelExecutableTask, RetryPolicy +from langgraph.types import PregelExecutableTask, RetryPolicy from langgraph.utils.config import patch_configurable logger = logging.getLogger(__name__) @@ -70,7 +70,9 @@ def run_with_retry( exc_info=exc, ) # signal subgraphs to resume (if available) - config = patch_configurable(config, {CONFIG_KEY_RESUMING: True}) + config = patch_configurable( + config, {CONFIG_KEY_RESUMING: True, CONFIG_KEY_GRAPH_COUNT: 0} + ) async def arun_with_retry( @@ -136,4 +138,6 @@ async def arun_with_retry( exc_info=exc, ) # signal subgraphs to resume (if available) - config = patch_configurable(config, {CONFIG_KEY_RESUMING: True}) + config = patch_configurable( + config, {CONFIG_KEY_RESUMING: True, CONFIG_KEY_GRAPH_COUNT: 0} + ) diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 6ba72e6ad..b8392b613 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -18,7 +18,7 @@ from langgraph.constants import ERROR, INTERRUPT, NO_WRITES from langgraph.errors import GraphDelegate, GraphInterrupt from langgraph.pregel.executor import Submit from langgraph.pregel.retry import arun_with_retry, run_with_retry -from langgraph.pregel.types import PregelExecutableTask, RetryPolicy +from langgraph.types import PregelExecutableTask, RetryPolicy class PregelRunner: diff --git a/libs/langgraph/langgraph/pregel/types.py b/libs/langgraph/langgraph/pregel/types.py index 4cca9946c..7a72b88c9 100644 --- a/libs/langgraph/langgraph/pregel/types.py +++ b/libs/langgraph/langgraph/pregel/types.py @@ -1,200 +1,25 @@ -from collections import deque -from dataclasses import dataclass -from typing import Any, Callable, Literal, NamedTuple, Optional, Sequence, Type, Union +"""Re-export types moved to langgraph.types""" -from langchain_core.runnables import Runnable, RunnableConfig +from langgraph.types import ( + All, + CachePolicy, + PregelExecutableTask, + PregelTask, + RetryPolicy, + StateSnapshot, + StreamMode, + StreamWriter, + default_retry_on, +) -from langgraph.checkpoint.base import CheckpointMetadata - -All = Literal["*"] - -StreamMode = Literal["values", "updates", "debug", "messages", "custom"] -"""How the stream method should emit outputs. - -- 'values': Emit all values of the state for each step. -- 'updates': Emit only the node name(s) and updates - that were returned by the node(s) **after** each step. -- 'debug': Emit debug events for each step. -- 'messages': Emit LLM messages token-by-token. -- 'custom': Emit custom output `write: StreamWriter` kwarg of each node. -""" - -StreamWriter = Callable[[Any], None] -"""Callable that accepts a single argument and writes it to the output stream. -Always injected into nodes if requested as a keyword argument, but it's a no-op -when not using stream_mode="custom".""" - - -def default_retry_on(exc: Exception) -> bool: - import httpx - import requests - - if isinstance(exc, ConnectionError): - return True - if isinstance( - exc, - ( - ValueError, - TypeError, - ArithmeticError, - ImportError, - LookupError, - NameError, - SyntaxError, - RuntimeError, - ReferenceError, - StopIteration, - StopAsyncIteration, - OSError, - ), - ): - return False - if isinstance(exc, httpx.HTTPStatusError): - return 500 <= exc.response.status_code < 600 - if isinstance(exc, requests.HTTPError): - return 500 <= exc.response.status_code < 600 if exc.response else True - return True - - -class RetryPolicy(NamedTuple): - """Configuration for retrying nodes.""" - - initial_interval: float = 0.5 - """Amount of time that must elapse before the first retry occurs. In seconds.""" - backoff_factor: float = 2.0 - """Multiplier by which the interval increases after each retry.""" - max_interval: float = 128.0 - """Maximum amount of time that may elapse between retries. In seconds.""" - max_attempts: int = 3 - """Maximum number of attempts to make before giving up, including the first.""" - jitter: bool = True - """Whether to add random jitter to the interval between retries.""" - retry_on: Union[ - Type[Exception], Sequence[Type[Exception]], Callable[[Exception], bool] - ] = default_retry_on - """List of exception classes that should trigger a retry, or a callable that returns True for exceptions that should trigger a retry.""" - - -class CachePolicy(NamedTuple): - """Configuration for caching nodes.""" - - pass - - -@dataclass -class Interrupt: - value: Any - when: Literal["during"] = "during" - - -class PregelTask(NamedTuple): - id: str - name: str - path: tuple[Union[str, int], ...] - error: Optional[Exception] = None - interrupts: tuple[Interrupt, ...] = () - state: Union[None, RunnableConfig, "StateSnapshot"] = None - - -class PregelExecutableTask(NamedTuple): - name: str - input: Any - proc: Runnable - writes: deque[tuple[str, Any]] - config: RunnableConfig - triggers: list[str] - retry_policy: Optional[RetryPolicy] - cache_policy: Optional[CachePolicy] - id: str - path: tuple[Union[str, int], ...] - scheduled: bool = False - - -class StateSnapshot(NamedTuple): - """Snapshot of the state of the graph at the beginning of a step.""" - - values: Union[dict[str, Any], Any] - """Current values of channels""" - next: tuple[str, ...] - """The name of the node to execute in each task for this step.""" - config: RunnableConfig - """Config used to fetch this snapshot""" - metadata: Optional[CheckpointMetadata] - """Metadata associated with this snapshot""" - created_at: Optional[str] - """Timestamp of snapshot creation""" - parent_config: Optional[RunnableConfig] - """Config used to fetch the parent snapshot, if any""" - tasks: tuple[PregelTask, ...] - """Tasks to execute in this step. If already attempted, may contain an error.""" - - -class Send: - """A message or packet to send to a specific node in the graph. - - The `Send` class is used within a `StateGraph`'s conditional edges to - dynamically invoke a node with a custom state at the next step. - - Importantly, the sent state can differ from the core graph's state, - allowing for flexible and dynamic workflow management. - - One such example is a "map-reduce" workflow where your graph invokes - the same node multiple times in parallel with different states, - before aggregating the results back into the main graph's state. - - Attributes: - node (str): The name of the target node to send the message to. - arg (Any): The state or message to send to the target node. - - Examples: - >>> from typing import Annotated - >>> import operator - >>> class OverallState(TypedDict): - ... subjects: list[str] - ... jokes: Annotated[list[str], operator.add] - ... - >>> from langgraph.constants import Send - >>> from langgraph.graph import END, START - >>> def continue_to_jokes(state: OverallState): - ... return [Send("generate_joke", {"subject": s}) for s in state['subjects']] - ... - >>> from langgraph.graph import StateGraph - >>> builder = StateGraph(OverallState) - >>> builder.add_node("generate_joke", lambda state: {"jokes": [f"Joke about {state['subject']}"]}) - >>> builder.add_conditional_edges(START, continue_to_jokes) - >>> builder.add_edge("generate_joke", END) - >>> graph = builder.compile() - >>> - >>> # Invoking with two subjects results in a generated joke for each - >>> graph.invoke({"subjects": ["cats", "dogs"]}) - {'subjects': ['cats', 'dogs'], 'jokes': ['Joke about cats', 'Joke about dogs']} - """ - - __slots__ = ("node", "arg") - - node: str - arg: Any - - def __init__(self, /, node: str, arg: Any) -> None: - """ - Initialize a new instance of the Send class. - - Args: - node (str): The name of the target node to send the message to. - arg (Any): The state or message to send to the target node. - """ - self.node = node - self.arg = arg - - def __hash__(self) -> int: - return hash((self.node, self.arg)) - - def __repr__(self) -> str: - return f"Send(node={self.node!r}, arg={self.arg!r})" - - def __eq__(self, value: object) -> bool: - return ( - isinstance(value, Send) - and self.node == value.node - and self.arg == value.arg - ) +__all__ = [ + "All", + "CachePolicy", + "PregelExecutableTask", + "PregelTask", + "RetryPolicy", + "StateSnapshot", + "StreamMode", + "StreamWriter", + "default_retry_on", +] diff --git a/libs/langgraph/langgraph/pregel/validate.py b/libs/langgraph/langgraph/pregel/validate.py index 965fab54e..cf957dc07 100644 --- a/libs/langgraph/langgraph/pregel/validate.py +++ b/libs/langgraph/langgraph/pregel/validate.py @@ -3,7 +3,7 @@ from typing import Any, Mapping, Optional, Sequence, Union from langgraph.channels.base import BaseChannel from langgraph.constants import RESERVED from langgraph.pregel.read import PregelNode -from langgraph.pregel.types import All +from langgraph.types import All def validate_graph( diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py new file mode 100644 index 000000000..4cca9946c --- /dev/null +++ b/libs/langgraph/langgraph/types.py @@ -0,0 +1,200 @@ +from collections import deque +from dataclasses import dataclass +from typing import Any, Callable, Literal, NamedTuple, Optional, Sequence, Type, Union + +from langchain_core.runnables import Runnable, RunnableConfig + +from langgraph.checkpoint.base import CheckpointMetadata + +All = Literal["*"] + +StreamMode = Literal["values", "updates", "debug", "messages", "custom"] +"""How the stream method should emit outputs. + +- 'values': Emit all values of the state for each step. +- 'updates': Emit only the node name(s) and updates + that were returned by the node(s) **after** each step. +- 'debug': Emit debug events for each step. +- 'messages': Emit LLM messages token-by-token. +- 'custom': Emit custom output `write: StreamWriter` kwarg of each node. +""" + +StreamWriter = Callable[[Any], None] +"""Callable that accepts a single argument and writes it to the output stream. +Always injected into nodes if requested as a keyword argument, but it's a no-op +when not using stream_mode="custom".""" + + +def default_retry_on(exc: Exception) -> bool: + import httpx + import requests + + if isinstance(exc, ConnectionError): + return True + if isinstance( + exc, + ( + ValueError, + TypeError, + ArithmeticError, + ImportError, + LookupError, + NameError, + SyntaxError, + RuntimeError, + ReferenceError, + StopIteration, + StopAsyncIteration, + OSError, + ), + ): + return False + if isinstance(exc, httpx.HTTPStatusError): + return 500 <= exc.response.status_code < 600 + if isinstance(exc, requests.HTTPError): + return 500 <= exc.response.status_code < 600 if exc.response else True + return True + + +class RetryPolicy(NamedTuple): + """Configuration for retrying nodes.""" + + initial_interval: float = 0.5 + """Amount of time that must elapse before the first retry occurs. In seconds.""" + backoff_factor: float = 2.0 + """Multiplier by which the interval increases after each retry.""" + max_interval: float = 128.0 + """Maximum amount of time that may elapse between retries. In seconds.""" + max_attempts: int = 3 + """Maximum number of attempts to make before giving up, including the first.""" + jitter: bool = True + """Whether to add random jitter to the interval between retries.""" + retry_on: Union[ + Type[Exception], Sequence[Type[Exception]], Callable[[Exception], bool] + ] = default_retry_on + """List of exception classes that should trigger a retry, or a callable that returns True for exceptions that should trigger a retry.""" + + +class CachePolicy(NamedTuple): + """Configuration for caching nodes.""" + + pass + + +@dataclass +class Interrupt: + value: Any + when: Literal["during"] = "during" + + +class PregelTask(NamedTuple): + id: str + name: str + path: tuple[Union[str, int], ...] + error: Optional[Exception] = None + interrupts: tuple[Interrupt, ...] = () + state: Union[None, RunnableConfig, "StateSnapshot"] = None + + +class PregelExecutableTask(NamedTuple): + name: str + input: Any + proc: Runnable + writes: deque[tuple[str, Any]] + config: RunnableConfig + triggers: list[str] + retry_policy: Optional[RetryPolicy] + cache_policy: Optional[CachePolicy] + id: str + path: tuple[Union[str, int], ...] + scheduled: bool = False + + +class StateSnapshot(NamedTuple): + """Snapshot of the state of the graph at the beginning of a step.""" + + values: Union[dict[str, Any], Any] + """Current values of channels""" + next: tuple[str, ...] + """The name of the node to execute in each task for this step.""" + config: RunnableConfig + """Config used to fetch this snapshot""" + metadata: Optional[CheckpointMetadata] + """Metadata associated with this snapshot""" + created_at: Optional[str] + """Timestamp of snapshot creation""" + parent_config: Optional[RunnableConfig] + """Config used to fetch the parent snapshot, if any""" + tasks: tuple[PregelTask, ...] + """Tasks to execute in this step. If already attempted, may contain an error.""" + + +class Send: + """A message or packet to send to a specific node in the graph. + + The `Send` class is used within a `StateGraph`'s conditional edges to + dynamically invoke a node with a custom state at the next step. + + Importantly, the sent state can differ from the core graph's state, + allowing for flexible and dynamic workflow management. + + One such example is a "map-reduce" workflow where your graph invokes + the same node multiple times in parallel with different states, + before aggregating the results back into the main graph's state. + + Attributes: + node (str): The name of the target node to send the message to. + arg (Any): The state or message to send to the target node. + + Examples: + >>> from typing import Annotated + >>> import operator + >>> class OverallState(TypedDict): + ... subjects: list[str] + ... jokes: Annotated[list[str], operator.add] + ... + >>> from langgraph.constants import Send + >>> from langgraph.graph import END, START + >>> def continue_to_jokes(state: OverallState): + ... return [Send("generate_joke", {"subject": s}) for s in state['subjects']] + ... + >>> from langgraph.graph import StateGraph + >>> builder = StateGraph(OverallState) + >>> builder.add_node("generate_joke", lambda state: {"jokes": [f"Joke about {state['subject']}"]}) + >>> builder.add_conditional_edges(START, continue_to_jokes) + >>> builder.add_edge("generate_joke", END) + >>> graph = builder.compile() + >>> + >>> # Invoking with two subjects results in a generated joke for each + >>> graph.invoke({"subjects": ["cats", "dogs"]}) + {'subjects': ['cats', 'dogs'], 'jokes': ['Joke about cats', 'Joke about dogs']} + """ + + __slots__ = ("node", "arg") + + node: str + arg: Any + + def __init__(self, /, node: str, arg: Any) -> None: + """ + Initialize a new instance of the Send class. + + Args: + node (str): The name of the target node to send the message to. + arg (Any): The state or message to send to the target node. + """ + self.node = node + self.arg = arg + + def __hash__(self) -> int: + return hash((self.node, self.arg)) + + def __repr__(self) -> str: + return f"Send(node={self.node!r}, arg={self.arg!r})" + + def __eq__(self, value: object) -> bool: + return ( + isinstance(value, Send) + and self.node == value.node + and self.arg == value.arg + ) diff --git a/libs/langgraph/langgraph/utils/runnable.py b/libs/langgraph/langgraph/utils/runnable.py index bfc2a627b..d81f86e34 100644 --- a/libs/langgraph/langgraph/utils/runnable.py +++ b/libs/langgraph/langgraph/utils/runnable.py @@ -35,7 +35,7 @@ from langchain_core.tracers._streaming import _StreamingCallbackHandler from typing_extensions import TypeGuard from langgraph.constants import CONFIG_KEY_STREAM_WRITER -from langgraph.pregel.types import StreamWriter +from langgraph.types import StreamWriter from langgraph.utils.config import ( ensure_config, get_async_callback_manager_for_config, diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 3e43ef28f..eb88ab189 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -70,8 +70,8 @@ from langgraph.pregel import ( StateSnapshot, ) from langgraph.pregel.retry import RetryPolicy -from langgraph.pregel.types import PregelTask, StreamWriter from langgraph.store.memory import MemoryStore +from langgraph.types import PregelTask, StreamWriter from tests.any_str import AnyDict, AnyStr, AnyVersion, UnsortedSequence from tests.conftest import ALL_CHECKPOINTERS_SYNC, SHOULD_CHECK_SNAPSHOTS from tests.fake_chat import FakeChatModel diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 640658399..cee63f469 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -68,8 +68,8 @@ from langgraph.pregel import ( StateSnapshot, ) from langgraph.pregel.retry import RetryPolicy -from langgraph.pregel.types import PregelTask, StreamWriter from langgraph.store.memory import MemoryStore +from langgraph.types import PregelTask, StreamWriter from tests.any_str import AnyDict, AnyStr, AnyVersion, UnsortedSequence from tests.conftest import ( ALL_CHECKPOINTERS_ASYNC, diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py index 9cbb6bf83..c803239e8 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py @@ -25,7 +25,6 @@ from langgraph.pregel.executor import ( ) from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager from langgraph.pregel.runner import PregelRunner -from langgraph.pregel.types import RetryPolicy from langgraph.scheduler.kafka.retry import aretry, retry from langgraph.scheduler.kafka.types import ( AsyncConsumer, @@ -38,6 +37,7 @@ from langgraph.scheduler.kafka.types import ( Sendable, Topics, ) +from langgraph.types import RetryPolicy from langgraph.utils.config import patch_configurable diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py index 1ad9c5c5b..097429bb6 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py @@ -24,7 +24,6 @@ from langgraph.errors import CheckpointNotLatest, GraphInterrupt from langgraph.pregel import Pregel from langgraph.pregel.executor import BackgroundExecutor, Submit from langgraph.pregel.loop import AsyncPregelLoop, SyncPregelLoop -from langgraph.pregel.types import RetryPolicy from langgraph.scheduler.kafka.retry import aretry, retry from langgraph.scheduler.kafka.types import ( AsyncConsumer, @@ -37,6 +36,7 @@ from langgraph.scheduler.kafka.types import ( Producer, Topics, ) +from langgraph.types import RetryPolicy from langgraph.utils.config import patch_configurable diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/retry.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/retry.py index bb80047f8..74dbe3e27 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/retry.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/retry.py @@ -6,7 +6,7 @@ from typing import Awaitable, Callable, Optional from typing_extensions import ParamSpec -from langgraph.pregel.types import RetryPolicy +from langgraph.types import RetryPolicy logger = logging.getLogger(__name__) P = ParamSpec("P")