cr

Dockerfile

@@ -1,4 +1,4 @@
-FROM python:3.9-slim
+FROM python:3.9
 
 # Set the working directory to /app
 WORKDIR /app
@@ -7,6 +7,6 @@ WORKDIR /app
 COPY . .
 
 # Install any needed packages specified in requirements.txt
-RUN pip install poetry && poetry config virtualenvs.create false && poetry install --with test,lint,typing
+RUN pip install poetry && poetry config virtualenvs.create false && poetry install --with test,lint,typing,dev
 
 RUN poetry run pytest

langgraph/graph/graph.py

@@ -1,6 +1,6 @@
 from asyncio import iscoroutinefunction
 from collections import defaultdict
-from typing import Any, Callable, Dict, NamedTuple
+from typing import Any, Callable, Dict, NamedTuple, Optional
 
 from langchain_core.runnables import Runnable
 from langchain_core.runnables.base import (
@@ -9,6 +9,7 @@ from langchain_core.runnables.base import (
     coerce_to_runnable,
 )
 
+from langgraph.checkpoint import BaseCheckpointSaver
 from langgraph.pregel import Channel, Pregel
 
 END = "__end__"
@@ -97,7 +98,7 @@ class Graph:
             if node not in all_starts:
                 raise ValueError(f"Node `{node}` is a dead-end")
 
-    def compile(self) -> Pregel:
+    def compile(self, checkpointer: Optional[BaseCheckpointSaver] = None) -> Pregel:
         self.validate()
 
         outgoing_edges = defaultdict(list)
@@ -127,4 +128,5 @@ class Graph:
             input=f"{self.entry_point}:inbox",
             output=END,
             hidden=[f"{node}:inbox" for node in self.nodes],
+            checkpointer=checkpointer,
         )

langgraph/graph/state.py

@@ -8,6 +8,7 @@ from langchain_core.runnables import RunnableConfig, RunnableLambda
 from langgraph.channels.base import BaseChannel
 from langgraph.channels.binop import BinaryOperatorAggregate
 from langgraph.channels.last_value import LastValue
+from langgraph.checkpoint import BaseCheckpointSaver
 from langgraph.graph.graph import END, Graph
 from langgraph.pregel import Channel, Pregel
 from langgraph.pregel.read import ChannelRead
@@ -24,7 +25,7 @@ class StateGraph(Graph):
         if any(isinstance(c, BinaryOperatorAggregate) for c in self.channels.values()):
             self.support_multiple_edges = True
 
-    def compile(self) -> Pregel:
+    def compile(self, checkpointer: Optional[BaseCheckpointSaver] = None) -> Pregel:
         self.validate()
 
         if any(key in self.nodes for key in self.channels):
@@ -79,6 +80,7 @@ class StateGraph(Graph):
             input=f"{START}:inbox",
             output=END,
             hidden=[f"{node}:inbox" for node in self.nodes] + [START] + state_keys,
+            checkpointer=checkpointer,
         )
 
 

langgraph/prebuilt/agent_executor.py

@@ -4,10 +4,17 @@ from typing import Annotated, TypedDict, Union, Sequence
 from langchain_core.agents import AgentAction, AgentFinish
 from langchain_core.messages import BaseMessage
 
+
+from typing import Annotated, Optional, TypedDict
+
+from langchain_core.agents import AgentAction, AgentFinish
+
+from langgraph.checkpoint import BaseCheckpointSaver
 from langgraph.graph import END, StateGraph
 from langgraph.prebuilt.tool_executor import ToolExecutor
 
 
+
 def _get_agent_state(input_schema= None):
     if input_schema is None:
         class AgentState(TypedDict):
@@ -36,8 +43,12 @@ def _get_agent_state(input_schema= None):
     return AgentState
 
 
-def create_agent_executor(agent_runnable, tools, input_schema=None):
-
+def create_agent_executor(
+    agent_runnable,
+    tools,
+    input_schema=None,
+    checkpointer: Optional[BaseCheckpointSaver] = None,
+):
     if isinstance(tools, ToolExecutor):
         tool_executor = tools
     else:
@@ -46,10 +57,11 @@ def create_agent_executor(agent_runnable, tools, input_schema=None):
     state = _get_agent_state(input_schema)
 
     # Define logic that will be used to determine which conditional edge to go down
+
     def should_continue(data):
         # If the agent outcome is an AgentFinish, then we return `exit` string
         # This will be used when setting up the graph to define the flow
-        if isinstance(data['agent_outcome'], AgentFinish):
+        if isinstance(data["agent_outcome"], AgentFinish):
             return "end"
         # Otherwise, an AgentAction is returned
         # Here we return `continue` string
@@ -64,7 +76,7 @@ def create_agent_executor(agent_runnable, tools, input_schema=None):
     # Define the function to execute tools
     def execute_tools(data):
         # Get the most recent agent_outcome - this is the key added in the `agent` above
-        agent_action = data['agent_outcome']
+        agent_action = data["agent_outcome"]
         output = tool_executor.invoke(agent_action)
         return {"intermediate_steps": [(agent_action, str(output))]}
 
@@ -96,15 +108,15 @@ def create_agent_executor(agent_runnable, tools, input_schema=None):
             # If `tools`, then we call the tool node.
             "continue": "action",
             # Otherwise we finish.
-            "end": END
-        }
+            "end": END,
+        },
     )
 
     # We now add a normal edge from `tools` to `agent`.
     # This means that after `tools` is called, `agent` node is called next.
-    workflow.add_edge('action', 'agent')
+    workflow.add_edge("action", "agent")
 
     # Finally, we compile it!
     # This compiles it into a LangChain Runnable,
     # meaning you can use it as you would any other runnable
-    return workflow.compile()
+    return workflow.compile(checkpointer=checkpointer)

langgraph/pregel/__init__.py

@@ -169,7 +169,7 @@ class Pregel(
 
     debug: bool = Field(default_factory=get_debug)
 
-    saver: Optional[BaseCheckpointSaver] = None
+    checkpointer: Optional[BaseCheckpointSaver] = None
 
     name: str = "LangGraph"
 
@@ -187,7 +187,7 @@ class Pregel(
     def config_specs(self) -> list[ConfigurableFieldSpec]:
         return get_unique_config_specs(
             [spec for node in self.nodes.values() for spec in node.config_specs]
-            + (self.saver.config_specs if self.saver is not None else [])
+            + (self.checkpointer.config_specs if self.checkpointer is not None else [])
         )
 
     @property
@@ -244,7 +244,7 @@ class Pregel(
         # copy nodes to ignore mutations during execution
         processes = {**self.nodes}
         # get checkpoint from saver, or create an empty one
-        checkpoint = self.saver.get(config) if self.saver else None
+        checkpoint = self.checkpointer.get(config) if self.checkpointer else None
         checkpoint = checkpoint or empty_checkpoint()
         # create channels from checkpoint
         with ChannelsManager(
@@ -327,18 +327,24 @@ class Pregel(
                         _apply_writes_from_view(checkpoint, channels, step_output)
 
                 # save end of step checkpoint
-                if self.saver is not None and self.saver.at == CheckpointAt.END_OF_STEP:
+                if (
+                    self.checkpointer is not None
+                    and self.checkpointer.at == CheckpointAt.END_OF_STEP
+                ):
                     checkpoint = create_checkpoint(checkpoint, channels)
-                    self.saver.put(config, checkpoint)
+                    self.checkpointer.put(config, checkpoint)
 
                 # interrupt if any channel written to is in interrupt list
                 if any(chan for chan, _ in pending_writes if chan in self.interrupt):
                     break
 
             # save end of run checkpoint
-            if self.saver is not None and self.saver.at == CheckpointAt.END_OF_RUN:
+            if (
+                self.checkpointer is not None
+                and self.checkpointer.at == CheckpointAt.END_OF_RUN
+            ):
                 checkpoint = create_checkpoint(checkpoint, channels)
-                self.saver.put(config, checkpoint)
+                self.checkpointer.put(config, checkpoint)
 
     async def _atransform(
         self,
@@ -368,7 +374,7 @@ class Pregel(
         # copy nodes to ignore mutations during execution
         processes = {**self.nodes}
         # get checkpoint from saver, or create an empty one
-        checkpoint = await self.saver.aget(config) if self.saver else None
+        checkpoint = await self.checkpointer.aget(config) if self.checkpointer else None
         checkpoint = checkpoint or empty_checkpoint()
         # create channels from checkpoint
         async with AsyncChannelsManager(self.channels, checkpoint) as channels:
@@ -454,18 +460,24 @@ class Pregel(
                         _apply_writes_from_view(checkpoint, channels, step_output)
 
                 # save end of step checkpoint
-                if self.saver is not None and self.saver.at == CheckpointAt.END_OF_STEP:
+                if (
+                    self.checkpointer is not None
+                    and self.checkpointer.at == CheckpointAt.END_OF_STEP
+                ):
                     checkpoint = create_checkpoint(checkpoint, channels)
-                    await self.saver.aput(config, checkpoint)
+                    await self.checkpointer.aput(config, checkpoint)
 
                 # interrupt if any channel written to is in interrupt list
                 if any(chan for chan, _ in pending_writes if chan in self.interrupt):
                     break
 
             # save end of run checkpoint
-            if self.saver is not None and self.saver.at == CheckpointAt.END_OF_RUN:
+            if (
+                self.checkpointer is not None
+                and self.checkpointer.at == CheckpointAt.END_OF_RUN
+            ):
                 checkpoint = create_checkpoint(checkpoint, channels)
-                await self.saver.aput(config, checkpoint)
+                await self.checkpointer.aput(config, checkpoint)
 
     def invoke(
         self,

tests/test_pregel.py

@@ -2,7 +2,7 @@ import operator
 import time
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import contextmanager
-from typing import Annotated, Generator, TypedDict
+from typing import Annotated, Generator, Optional, TypedDict, Union
 
 import pytest
 from langchain_core.runnables import RunnablePassthrough
@@ -382,7 +382,7 @@ def test_invoke_checkpoint(mocker: MockerFixture) -> None:
     app = Pregel(
         nodes={"one": one},
         channels={"total": BinaryOperatorAggregate(int, operator.add)},
-        saver=memory,
+        checkpointer=memory,
     )
 
     # total starts out as 0, so output is 0+2=2
@@ -584,7 +584,7 @@ def test_conditional_graph() -> None:
         ]
     )
 
-    def agent_parser(input: str) -> AgentFinish | AgentAction:
+    def agent_parser(input: str) -> Union[AgentAction, AgentFinish]:
         if input.startswith("finish"):
             _, answer = input.split(":")
             return AgentFinish(return_values={"answer": answer}, log=input)
@@ -786,7 +786,7 @@ def test_conditional_graph_state() -> None:
 
     class AgentState(TypedDict):
         input: str
-        agent_outcome: AgentAction | AgentFinish | None
+        agent_outcome: Optional[Union[AgentAction, AgentFinish]]
         intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add]
 
     # Assemble the tools
@@ -808,7 +808,7 @@ def test_conditional_graph_state() -> None:
         ]
     )
 
-    def agent_parser(input: str) -> AgentFinish | AgentAction:
+    def agent_parser(input: str) -> Union[AgentAction, AgentFinish]:
         if input.startswith("finish"):
             _, answer = input.split(":")
             return {

tests/test_pregel_async.py

@@ -1,7 +1,16 @@
 import asyncio
 import operator
 from contextlib import asynccontextmanager, contextmanager
-from typing import Annotated, Any, AsyncGenerator, AsyncIterator, Generator, TypedDict
+from typing import (
+    Annotated,
+    Any,
+    AsyncGenerator,
+    AsyncIterator,
+    Generator,
+    Optional,
+    TypedDict,
+    Union,
+)
 
 import pytest
 from langchain_core.runnables import RunnablePassthrough
@@ -396,7 +405,7 @@ async def test_invoke_checkpoint(mocker: MockerFixture) -> None:
     app = Pregel(
         nodes={"one": one},
         channels={"total": BinaryOperatorAggregate(int, operator.add)},
-        saver=memory,
+        checkpointer=memory,
     )
 
     # total starts out as 0, so output is 0+2=2
@@ -621,7 +630,7 @@ async def test_conditional_graph() -> None:
         ]
     )
 
-    async def agent_parser(input: str) -> AgentFinish | AgentAction:
+    async def agent_parser(input: str) -> Union[AgentAction, AgentFinish]:
         if input.startswith("finish"):
             _, answer = input.split(":")
             return AgentFinish(return_values={"answer": answer}, log=input)
@@ -831,7 +840,7 @@ async def test_conditional_graph_state() -> None:
 
     class AgentState(TypedDict):
         input: str
-        agent_outcome: AgentAction | AgentFinish | None
+        agent_outcome: Optional[Union[AgentAction, AgentFinish]]
         intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add]
 
     # Assemble the tools
@@ -853,7 +862,7 @@ async def test_conditional_graph_state() -> None:
         ]
     )
 
-    def agent_parser(input: str) -> AgentFinish | AgentAction:
+    def agent_parser(input: str) -> Union[AgentAction, AgentFinish]:
         if input.startswith("finish"):
             _, answer = input.split(":")
             return {