diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py index a5f657824..31bffd19c 100644 --- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py +++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py @@ -165,14 +165,15 @@ class SqliteSaver(BaseCheckpointSaver): Yields: sqlite3.Cursor: A cursor for the SQLite database. """ - self.setup() - cur = self.conn.cursor() - try: - yield cur - finally: - if transaction: - self.conn.commit() - cur.close() + with self.lock: + self.setup() + cur = self.conn.cursor() + try: + yield cur + finally: + if transaction: + self.conn.commit() + cur.close() def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: """Get a checkpoint tuple from the database. @@ -398,7 +399,7 @@ class SqliteSaver(BaseCheckpointSaver): checkpoint_ns = config["configurable"]["checkpoint_ns"] type_, serialized_checkpoint = self.serde.dumps_typed(checkpoint) serialized_metadata = self.jsonplus_serde.dumps(metadata) - with self.lock, self.cursor() as cur: + with self.cursor() as cur: cur.execute( "INSERT OR REPLACE INTO checkpoints (thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata) VALUES (?, ?, ?, ?, ?, ?, ?)", ( @@ -434,7 +435,7 @@ class SqliteSaver(BaseCheckpointSaver): writes (Sequence[Tuple[str, Any]]): List of writes to store, each as (channel, value) pair. task_id (str): Identifier for the task creating the writes. """ - with self.lock, self.cursor() as cur: + with self.cursor() as cur: cur.executemany( "INSERT OR IGNORE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)", [ diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py index 7cdc7c8fb..28ad76109 100644 --- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py +++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py @@ -276,7 +276,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver): """ await self.setup() checkpoint_ns = config["configurable"].get("checkpoint_ns", "") - async with self.conn.cursor() as cur: + async with self.lock, self.conn.cursor() as cur: # find the latest checkpoint for the thread_id if checkpoint_id := get_checkpoint_id(config): await cur.execute( @@ -371,7 +371,9 @@ class AsyncSqliteSaver(BaseCheckpointSaver): ORDER BY checkpoint_id DESC""" if limit: query += f" LIMIT {limit}" - async with self.conn.execute(query, params) as cur, self.conn.cursor() as wcur: + async with self.lock, self.conn.execute( + query, params + ) as cur, self.conn.cursor() as wcur: async for ( thread_id, checkpoint_ns, @@ -438,7 +440,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver): checkpoint_ns = config["configurable"]["checkpoint_ns"] type_, serialized_checkpoint = self.serde.dumps_typed(checkpoint) serialized_metadata = self.jsonplus_serde.dumps(metadata) - async with self.conn.execute( + async with self.lock, self.conn.execute( "INSERT OR REPLACE INTO checkpoints (thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata) VALUES (?, ?, ?, ?, ?, ?, ?)", ( str(config["configurable"]["thread_id"]), @@ -475,7 +477,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver): task_id (str): Identifier for the task creating the writes. """ await self.setup() - async with self.conn.cursor() as cur: + async with self.lock, self.conn.cursor() as cur: await cur.executemany( "INSERT OR IGNORE INTO writes (thread_id, checkpoint_ns, checkpoint_id, task_id, idx, channel, type, value) VALUES (?, ?, ?, ?, ?, ?, ?, ?)", [ diff --git a/libs/langgraph/langgraph/graph/graph.py b/libs/langgraph/langgraph/graph/graph.py index d466a3c96..570c06bad 100644 --- a/libs/langgraph/langgraph/graph/graph.py +++ b/libs/langgraph/langgraph/graph/graph.py @@ -424,6 +424,10 @@ class Graph: class CompiledGraph(Pregel): builder: Graph + def __init__(self, *, builder: Graph, **kwargs): + super().__init__(**kwargs) + self.builder = builder + def attach_node(self, key: str, node: NodeSpec) -> None: self.channels[key] = EphemeralValue(Any) self.nodes[key] = ( diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index f576a98cd..307b672aa 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -17,10 +17,10 @@ from typing import ( overload, ) -from langchain_core.pydantic_v1 import BaseModel from langchain_core.runnables import Runnable, RunnableConfig from langchain_core.runnables.base import RunnableLike from langchain_core.runnables.utils import create_model +from pydantic import BaseModel from langgraph.channels.base import BaseChannel from langgraph.channels.binop import BinaryOperatorAggregate diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 01f0ad91d..d7c265b2d 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -25,12 +25,10 @@ from uuid import UUID, uuid5 from langchain_core.globals import get_debug from langchain_core.load.dump import dumpd -from langchain_core.pydantic_v1 import BaseModel, Field, root_validator from langchain_core.runnables import ( Runnable, RunnableLambda, RunnableSequence, - RunnableSerializable, ) from langchain_core.runnables.base import Input, Output, coerce_to_runnable from langchain_core.runnables.config import ( @@ -48,6 +46,7 @@ from langchain_core.runnables.utils import ( get_unique_config_specs, ) from langchain_core.tracers._streaming import _StreamingCallbackHandler +from pydantic import BaseModel from typing_extensions import Self from langgraph.channels.base import ( @@ -186,16 +185,10 @@ class Channel: ) -class Pregel( - RunnableSerializable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]] -): +class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]): nodes: Mapping[str, PregelNode] - channels: Mapping[str, Union[BaseChannel, ManagedValueSpec]] = Field( - default_factory=dict - ) - - auto_validate: bool = True + channels: Mapping[str, Union[BaseChannel, ManagedValueSpec]] stream_mode: StreamMode = "values" """Mode to stream output, defaults to 'values'.""" @@ -205,16 +198,16 @@ class Pregel( stream_channels: Optional[Union[str, Sequence[str]]] = None """Channels to stream, defaults to all channels not in reserved channels""" - interrupt_after_nodes: Union[All, Sequence[str]] = Field(default_factory=list) + interrupt_after_nodes: Union[All, Sequence[str]] - interrupt_before_nodes: Union[All, Sequence[str]] = Field(default_factory=list) + interrupt_before_nodes: Union[All, Sequence[str]] input_channels: Union[str, Sequence[str]] step_timeout: Optional[float] = None """Maximum time to wait for a step to complete, in seconds. Defaults to None.""" - debug: bool = Field(default_factory=get_debug) + debug: bool """Whether to print debug information during execution. Defaults to False.""" checkpointer: Optional[BaseCheckpointSaver] = None @@ -232,36 +225,50 @@ class Pregel( name: str = "LangGraph" - class Config: - arbitrary_types_allowed = True + def __init__( + self, + *, + nodes: Mapping[str, PregelNode], + channels: Mapping[str, Union[BaseChannel, ManagedValueSpec]] = None, + auto_validate: bool = True, + stream_mode: StreamMode = "values", + output_channels: Union[str, Sequence[str]], + stream_channels: Optional[Union[str, Sequence[str]]] = None, + interrupt_after_nodes: Union[All, Sequence[str]] = (), + interrupt_before_nodes: Union[All, Sequence[str]] = (), + input_channels: Union[str, Sequence[str]], + step_timeout: Optional[float] = None, + debug: Optional[bool] = None, + checkpointer: Optional[BaseCheckpointSaver] = None, + store: Optional[BaseStore] = None, + retry_policy: Optional[RetryPolicy] = None, + config_type: Optional[Type[Any]] = None, + config: Optional[RunnableConfig] = None, + name: str = "LangGraph", + ) -> None: + self.nodes = nodes + self.channels = channels or {} + self.stream_mode = stream_mode + self.output_channels = output_channels + self.stream_channels = stream_channels + self.interrupt_after_nodes = interrupt_after_nodes + self.interrupt_before_nodes = interrupt_before_nodes + self.input_channels = input_channels + self.step_timeout = step_timeout + self.debug = debug if debug is not None else get_debug() + self.checkpointer = checkpointer + self.store = store + self.retry_policy = retry_policy + self.config_type = config_type + self.config = config + self.name = name + if auto_validate: + self.validate() def with_config(self, config: RunnableConfig | None = None, **kwargs: Any) -> Self: - return self.copy( - update={"config": cast(RunnableConfig, {**(config or {}), **kwargs})} - ) - - @classmethod - def is_lc_serializable(cls) -> bool: - """Return whether the graph can be serialized by Langchain.""" - return True - - @root_validator(skip_on_failure=True) - def validate_on_init(cls, values: dict[str, Any]) -> dict[str, Any]: - if not values["auto_validate"]: - return values - validate_graph( - values["nodes"], - values["channels"], - values["input_channels"], - values["output_channels"], - values["stream_channels"], - values["interrupt_after_nodes"], - values["interrupt_before_nodes"], - ) - if values["interrupt_after_nodes"] or values["interrupt_before_nodes"]: - if not values["checkpointer"]: - raise ValueError("Interrupts require a checkpointer") - return values + attrs = {**self.__dict__} + attrs["config"] = merge_configs(self.config, config, kwargs) + return self.__class__(**attrs) def validate(self) -> Self: validate_graph( diff --git a/libs/langgraph/langgraph/pregel/read.py b/libs/langgraph/langgraph/pregel/read.py index 163665e2e..ca0828fda 100644 --- a/libs/langgraph/langgraph/pregel/read.py +++ b/libs/langgraph/langgraph/pregel/read.py @@ -1,8 +1,16 @@ from __future__ import annotations -from typing import Any, Callable, Mapping, Optional, Sequence, Union +from typing import ( + Any, + AsyncIterator, + Callable, + Iterator, + Mapping, + Optional, + Sequence, + Union, +) -from langchain_core.pydantic_v1 import Field from langchain_core.runnables import ( Runnable, RunnableConfig, @@ -10,7 +18,7 @@ from langchain_core.runnables import ( RunnableSequence, RunnableSerializable, ) -from langchain_core.runnables.base import Other, RunnableBindingBase, coerce_to_runnable +from langchain_core.runnables.base import Input, Other, Output, coerce_to_runnable from langchain_core.runnables.config import merge_configs from langchain_core.runnables.utils import ConfigurableFieldSpec @@ -99,20 +107,47 @@ class ChannelRead(RunnableCallable): DEFAULT_BOUND: RunnablePassthrough = RunnablePassthrough() -class PregelNode(RunnableBindingBase): +class PregelNode(Runnable): channels: Union[list[str], Mapping[str, str]] - triggers: list[str] = Field(default_factory=list) + triggers: list[str] - mapper: Optional[Callable[[Any], Any]] = None + mapper: Optional[Callable[[Any], Any]] - writers: list[Runnable] = Field(default_factory=list) + writers: list[Runnable] - bound: Runnable[Any, Any] = Field(default=DEFAULT_BOUND) + bound: Runnable[Any, Any] - kwargs: Mapping[str, Any] = Field(default_factory=dict) + retry_policy: Optional[RetryPolicy] - retry_policy: Optional[RetryPolicy] = None + config: RunnableConfig + + def __init__( + self, + *, + channels: Union[list[str], Mapping[str, str]], + triggers: Sequence[str], + mapper: Optional[Callable[[Any], Any]] = None, + writers: Optional[list[Runnable]] = None, + tags: Optional[list[str]] = None, + metadata: Optional[Mapping[str, Any]] = None, + bound: Optional[Runnable[Any, Any]] = None, + retry_policy: Optional[RetryPolicy] = None, + config: Optional[RunnableConfig] = None, + ) -> None: + self.channels = channels + self.triggers = list(triggers) + self.mapper = mapper + self.writers = writers or [] + self.bound = bound if bound is not None else DEFAULT_BOUND + self.retry_policy = retry_policy + self.config = merge_configs( + config, {"tags": tags or [], "metadata": metadata or {}} + ) + + def copy(self, update: dict[str, Any]) -> PregelNode: + attrs = {**self.__dict__, **update} + return PregelNode(**attrs) def get_writers(self) -> list[Runnable]: """Get writers with optimizations applied.""" @@ -145,38 +180,6 @@ class PregelNode(RunnableBindingBase): else: return self.bound - def __init__( - self, - *, - channels: Union[list[str], Mapping[str, str]], - triggers: Sequence[str], - mapper: Optional[Callable[[Any], Any]] = None, - writers: Optional[list[Runnable]] = None, - tags: Optional[list[str]] = None, - metadata: Optional[Mapping[str, Any]] = None, - bound: Optional[Runnable[Any, Any]] = None, - kwargs: Optional[Mapping[str, Any]] = None, - config: Optional[RunnableConfig] = None, - retry_policy: Optional[RetryPolicy] = None, - **other_kwargs: Any, - ) -> None: - super().__init__( - channels=channels, - triggers=triggers, - mapper=mapper, - writers=writers or [], - bound=bound or DEFAULT_BOUND, - kwargs=kwargs or {}, - retry_policy=retry_policy, - config=merge_configs( - config, {"tags": tags or [], "metadata": metadata or {}} - ), - **other_kwargs, - ) - - def __repr_args__(self) -> Any: - return [(k, v) for k, v in super().__repr_args__() if k != "bound"] - def join(self, channels: Sequence[str]) -> PregelNode: assert isinstance(channels, list) or isinstance( channels, tuple @@ -226,3 +229,42 @@ class PregelNode(RunnableBindingBase): ], ) -> RunnableSerializable: raise NotImplementedError() + + def invoke( + self, + input: Input, + config: Optional[RunnableConfig] = None, + **kwargs: Optional[Any], + ) -> Output: + return self.bound.invoke(input, merge_configs(self.config, config), **kwargs) + + async def ainvoke( + self, + input: Input, + config: Optional[RunnableConfig] = None, + **kwargs: Optional[Any], + ) -> Output: + return await self.bound.ainvoke( + input, merge_configs(self.config, config), **kwargs + ) + + def stream( + self, + input: Input, + config: Optional[RunnableConfig] = None, + **kwargs: Optional[Any], + ) -> Iterator[Output]: + yield from self.bound.stream( + input, merge_configs(self.config, config), **kwargs + ) + + async def astream( + self, + input: Input, + config: Optional[RunnableConfig] = None, + **kwargs: Optional[Any], + ) -> AsyncIterator[Output]: + async for item in self.bound.astream( + input, merge_configs(self.config, config), **kwargs + ): + yield item diff --git a/libs/langgraph/tests/__snapshots__/test_pregel.ambr b/libs/langgraph/tests/__snapshots__/test_pregel.ambr index 3c3dd49d0..3937546f5 100644 --- a/libs/langgraph/tests/__snapshots__/test_pregel.ambr +++ b/libs/langgraph/tests/__snapshots__/test_pregel.ambr @@ -3304,7 +3304,7 @@ 'query', 'inner', ]), - 'title': 'Input', + 'title': 'LangGraphInput', 'type': 'object', }) # --- @@ -3323,11 +3323,7 @@ 'type': 'array', }), }), - 'required': list([ - 'answer', - 'docs', - ]), - 'title': 'Output', + 'title': 'LangGraphOutput', 'type': 'object', }) # --- @@ -3374,7 +3370,7 @@ 'query', 'inner', ]), - 'title': 'Input', + 'title': 'LangGraphInput', 'type': 'object', }) # --- @@ -3393,11 +3389,7 @@ 'type': 'array', }), }), - 'required': list([ - 'answer', - 'docs', - ]), - 'title': 'Output', + 'title': 'LangGraphOutput', 'type': 'object', }) # --- @@ -3444,7 +3436,7 @@ 'query', 'inner', ]), - 'title': 'Input', + 'title': 'LangGraphInput', 'type': 'object', }) # --- @@ -3463,11 +3455,7 @@ 'type': 'array', }), }), - 'required': list([ - 'answer', - 'docs', - ]), - 'title': 'Output', + 'title': 'LangGraphOutput', 'type': 'object', }) # --- @@ -3514,7 +3502,7 @@ 'query', 'inner', ]), - 'title': 'Input', + 'title': 'LangGraphInput', 'type': 'object', }) # --- @@ -3533,11 +3521,7 @@ 'type': 'array', }), }), - 'required': list([ - 'answer', - 'docs', - ]), - 'title': 'Output', + 'title': 'LangGraphOutput', 'type': 'object', }) # --- @@ -3584,7 +3568,7 @@ 'query', 'inner', ]), - 'title': 'Input', + 'title': 'LangGraphInput', 'type': 'object', }) # --- @@ -3603,11 +3587,7 @@ 'type': 'array', }), }), - 'required': list([ - 'answer', - 'docs', - ]), - 'title': 'Output', + 'title': 'LangGraphOutput', 'type': 'object', }) # --- diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 3ca3906d1..96696fa13 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -1903,7 +1903,7 @@ def test_invoke_two_processes_no_in(mocker: MockerFixture) -> None: one = Channel.subscribe_to("between") | add_one | Channel.write_to("output") two = Channel.subscribe_to("between") | add_one - with pytest.raises(ValueError): + with pytest.raises(TypeError): Pregel(nodes={"one": one, "two": two}) @@ -9943,10 +9943,12 @@ def test_send_to_nested_graphs( graph.update_state(outer_state.tasks[1].state, {"subject": "turtles - hohoho"}) # continue past interrupt - assert graph.invoke(None, config=config) == { - "subjects": ["cats", "dogs"], - "jokes": ["Joke about cats - hohoho", "Joke about turtles - hohoho"], - } + assert sorted( + graph.stream(None, config=config), key=lambda d: d["generate_joke"]["jokes"][0] + ) == [ + {"generate_joke": {"jokes": ["Joke about cats - hohoho"]}}, + {"generate_joke": {"jokes": ["Joke about turtles - hohoho"]}}, + ] actual_snapshot = graph.get_state(config) expected_snapshot = StateSnapshot(