WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-22 15:42:25 +02:00
Code Issues Packages Projects Releases Wiki Activity

Rename saver to checkpointer, expose in graph, state graph, prebuilt agent exec

Browse Source
This commit is contained in:
Nuno Campos
2024-01-15 13:45:54 -08:00
parent 23f84c9ac9
commit 2e537320c3
6 changed files with 52 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
langgraph/graph/graph.py
+4 -2
View File
@@ -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
+3 -1
View File
@@ -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
+19 -13
View File
@@ -1,27 +1,33 @@
from typing import Annotated, TypedDict
import operator
from typing import Annotated, Optional, TypedDict
from langchain_core.agents import AgentAction, AgentFinish
from langgraph.graph import StateGraph, END
from langgraph.checkpoint import BaseCheckpointSaver
from langgraph.graph import END, StateGraph
from langgraph.prebuilt.tool_executor import ToolExecutor
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:
tool_executor = ToolExecutor(tools)
if input_schema is None:
class AgentState(TypedDict):
input: str
agent_outcome: AgentAction | AgentFinish | None
intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add]
else:
class AgentState(input_schema):
agent_outcome: AgentAction | AgentFinish | None
intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add]
@@ -29,7 +35,7 @@ def create_agent_executor(agent_runnable, tools, input_schema=None):
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
@@ -44,7 +50,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))]}
@@ -76,15 +82,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
+24 -12
View File
@@ -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
+1 -1
View File
@@ -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
tests/test_pregel_async.py
+1 -1
View File
@@ -405,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