diff --git a/.github/workflows/deploy_docs.yml b/.github/workflows/deploy_docs.yml index 84de2e178..e051d3750 100644 --- a/.github/workflows/deploy_docs.yml +++ b/.github/workflows/deploy_docs.yml @@ -85,7 +85,8 @@ jobs: if [ "${{ github.event_name }}" == "schedule" ] || [ "${{ github.event_name }}" == "workflow_dispatch" ] || ([ "${{ github.event_name }}" == "push" ] && [ "${{ github.ref }}" == "refs/heads/main" ]); then echo "Running link check on all HTML files matching notebooks in docs directory..." poetry run pytest -v \ - --check-links-ignore "https://(api|web|docs|academy)\.smith\.langchain\.com/.*" \ + --check-links-ignore "https://(api|web|docs)\.smith\.langchain\.com/.*" \ + --check-links-ignore "https://academy\.langchain\.com/.*" \ --check-links-ignore "https://x.com/.*" \ --check-links-ignore "https://github\.com/.*" \ --check-links-ignore "http://localhost:8123/.*" \ @@ -106,7 +107,8 @@ jobs: if [ -n "${CHANGED_FILES}" ]; then echo "Running link check on HTML files matching changed notebook files..." poetry run pytest -v \ - --check-links-ignore "https://(api|web|docs|academy)\.smith\.langchain\.com/.*" \ + --check-links-ignore "https://(api|web|docs)\.smith\.langchain\.com/.*" \ + --check-links-ignore "https://academy\.langchain\.com/.*" \ --check-links-ignore "http://localhost:8123/.*" \ --check-links-ignore "http://localhost:2024.*" \ --check-links-ignore "http://127.0.0.1:.*" \ diff --git a/docs/docs/cloud/reference/api/openapi_control_plane.json b/docs/docs/cloud/reference/api/openapi_control_plane.json index 646ca09b6..871291505 100644 --- a/docs/docs/cloud/reference/api/openapi_control_plane.json +++ b/docs/docs/cloud/reference/api/openapi_control_plane.json @@ -561,6 +561,16 @@ }, "resource": { "$ref": "#/components/schemas/ResourceService" + }, + "status": { + "type": "string", + "enum": [ + "AWAITING_DATABASE", + "READY", + "AWAITING_DELETE", + "UNKNOWN" + ], + "description": "Deployment status of the project.\n\nNon-terminal statuses: `AWAITING_DATABASE`, `AWAITING_DELETE`. All other statuses are terminal." } } }, diff --git a/docs/docs/concepts/langgraph_cloud.md b/docs/docs/concepts/langgraph_cloud.md index 169953b8b..6d32d74a1 100644 --- a/docs/docs/concepts/langgraph_cloud.md +++ b/docs/docs/concepts/langgraph_cloud.md @@ -6,21 +6,29 @@ ## Overview -LangGraph's Cloud SaaS is a managed service for deploying LangGraph APIs, regardless of its definition or dependencies. The service offers managed implementations of checkpointers and stores, allowing you to focus on building the right cognitive architecture for your use case. By handling scalable & secure infrastructure, LangGraph Cloud offers the fastest path to getting your LangGraph API deployed to production. +LangGraph's Cloud SaaS is a managed service for deploying LangGraph Servers, regardless of its definition or dependencies. The service offers managed implementations of checkpointers and stores, allowing you to focus on building the right cognitive architecture for your use case. By handling scalable & secure infrastructure, LangGraph Cloud SaaS offers the fastest path to getting your LangGraph Server deployed to production. ## Deployment -A **deployment** is an instance of a LangGraph API. A single deployment can have many [revisions](#revision). When a deployment is created, all the necessary infrastructure (e.g. database, containers, secrets store) are automatically provisioned. See the [architecture diagram](#architecture) below for more details. +A **deployment** is an instance of a LangGraph Server. A single deployment can have many [revisions](#revision). When a deployment is created, all the necessary infrastructure (e.g. database, containers, secrets store) are automatically provisioned. See the [architecture diagram](#architecture) below for more details. -See the [how-to guide](../cloud/deployment/cloud.md#create-new-deployment) for creating a new deployment. - -## Resource Allocation +Resource Allocation: | **Deployment Type** | **CPU** | **Memory** | **Scaling** | |---------------------|---------|------------|---------------------| | Development | 1 CPU | 1 GB | Up to 1 container | | Production | 2 CPU | 2 GB | Up to 10 containers | +See the [how-to guide](../cloud/deployment/cloud.md#create-new-deployment) for creating a new deployment. + +## Persistence + +A dedicated database is automatically created for each deployment. The database serves as the [persistence layer](../concepts/persistence.md) for the deployment. + +When defining a graph to be deployed to LangGraph Cloud SaaS, a [checkpointer](../concepts/persistence.md#checkpointer-libraries) should not be configured by the user. Instead, a checkpointer is automatically configured for the graph. + +There is no direct access to the database. All access to the database occurs through the LangGraph Server APIs. + ## Autoscaling `Production` type deployments automatically scale up to 10 containers. Scaling is based on the current request load for a single container. Specifically, the autoscaling implementation scales the deployment so that each container is processing about 10 concurrent requests. For example... diff --git a/docs/mkdocs.yml b/docs/mkdocs.yml index a6e159f15..2a5079ae4 100644 --- a/docs/mkdocs.yml +++ b/docs/mkdocs.yml @@ -20,6 +20,7 @@ theme: - content.action.edit - content.tooltips - header.autohide + - navigation.indexes - navigation.expand - navigation.footer - navigation.instant diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py index 89275f111..0d901a78c 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py @@ -57,6 +57,9 @@ MIGRATIONS = [ PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx) );""", "ALTER TABLE checkpoint_blobs ALTER COLUMN blob DROP not null;", + # NOTE: this is a no-op migration to ensure that the versions in the migrations table are correct. + # This is necessary due to an empty migration previously added to the list. + "SELECT 1;", """ CREATE INDEX CONCURRENTLY IF NOT EXISTS checkpoints_thread_id_idx ON checkpoints(thread_id); """, diff --git a/libs/checkpoint-postgres/pyproject.toml b/libs/checkpoint-postgres/pyproject.toml index addcdfe20..4a62d75fc 100644 --- a/libs/checkpoint-postgres/pyproject.toml +++ b/libs/checkpoint-postgres/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph-checkpoint-postgres" -version = "2.0.12" +version = "2.0.13" description = "Library with a Postgres implementation of LangGraph checkpoint saver." authors = [] license = "MIT" diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index cb6834f5f..fd5f8758c 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -75,11 +75,11 @@ CONFIG_KEY_CHECKPOINT_ID = sys.intern("checkpoint_id") CONFIG_KEY_CHECKPOINT_NS = sys.intern("checkpoint_ns") # holds the current checkpoint_ns, "" for root graph CONFIG_KEY_NODE_FINISHED = sys.intern("__pregel_node_finished") -# holds the value that "answers" an interrupt() call -CONFIG_KEY_WRITES = sys.intern("__pregel_writes") -# read-only list of existing task writes +# holds a callback to be called when a node is finished CONFIG_KEY_SCRATCHPAD = sys.intern("__pregel_scratchpad") # holds a mutable dict for temporary storage scoped to the current task +CONFIG_KEY_END = sys.intern("__pregel_previous") +# holds the previous return value from a stateful Pregel graph. # --- Other constants --- PUSH = sys.intern("__pregel_push") diff --git a/libs/langgraph/langgraph/errors.py b/libs/langgraph/langgraph/errors.py index 0737a31d0..8e78a8784 100644 --- a/libs/langgraph/langgraph/errors.py +++ b/libs/langgraph/langgraph/errors.py @@ -107,18 +107,3 @@ class CheckpointNotLatest(Exception): """Raised when the checkpoint is not the latest version (for distributed mode).""" pass - - -class MultipleSubgraphsError(Exception): - """Raised when multiple subgraphs are called inside the same node. - - Troubleshooting guides: - - - [MULTIPLE_SUBGRAPHS](https://python.langchain.com/docs/troubleshooting/errors/MULTIPLE_SUBGRAPHS) - """ - - pass - - -_SEEN_CHECKPOINT_NS: set[str] = set() -"""Used for subgraph detection.""" diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index d625be560..a051e465e 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -110,23 +110,99 @@ def entrypoint( *, checkpointer: Optional[BaseCheckpointSaver] = None, store: Optional[BaseStore] = None, + config_schema: Optional[type[Any]] = None, ) -> Callable[[types.FunctionType], Pregel]: def _imp(func: types.FunctionType) -> Pregel: + """Convert a function into a Pregel graph. + + Args: + func: The function to convert. Support both sync and async functions, as well + as generator and async generator functions. + + Returns: + A Pregel graph. + """ + # wrap generators in a function that writes to StreamWriter if inspect.isgeneratorfunction(func): + original_sig = inspect.signature(func) + # Check if original signature has a writer argument with a matching type. + # If not, we'll inject it into the decorator, but not pass it + # to the wrapped function. + if "writer" in original_sig.parameters: - def gen_wrapper(*args: Any, writer: StreamWriter, **kwargs: Any) -> Any: - for chunk in func(*args, **kwargs): - writer(chunk) + @functools.wraps(func) + def gen_wrapper(*args: Any, writer: StreamWriter, **kwargs: Any) -> Any: + chunks = [] + for chunk in func(*args, writer=writer, **kwargs): + writer(chunk) + chunks.append(chunk) + return chunks + else: + @functools.wraps(func) + def gen_wrapper(*args: Any, writer: StreamWriter, **kwargs: Any) -> Any: + chunks = [] + # Do not pass the writer argument to the wrapped function + # as it does not have a matching parameter + for chunk in func(*args, **kwargs): + writer(chunk) + chunks.append(chunk) + return chunks + + # Create a new parameter for the writer argument + extra_param = inspect.Parameter( + "writer", + inspect.Parameter.KEYWORD_ONLY, + # The extra argument is a keyword-only argument + default=lambda _: None, + ) + # Update the function's signature to include the extra argument + new_params = list(original_sig.parameters.values()) + [extra_param] + new_sig = original_sig.replace(parameters=new_params) + # Update the signature of the wrapper function + gen_wrapper.__signature__ = new_sig # type: ignore bound = get_runnable_for_func(gen_wrapper) stream_mode: StreamMode = "custom" elif inspect.isasyncgenfunction(func): + original_sig = inspect.signature(func) + # Check if original signature has a writer argument with a matching type. + # If not, we'll inject it into the decorator, but not pass it + # to the wrapped function. + if "writer" in original_sig.parameters: - async def agen_wrapper( - *args: Any, writer: StreamWriter, **kwargs: Any - ) -> Any: - async for chunk in func(*args, **kwargs): - writer(chunk) + @functools.wraps(func) + async def agen_wrapper( + *args: Any, writer: StreamWriter, **kwargs: Any + ) -> Any: + chunks = [] + async for chunk in func(*args, writer=writer, **kwargs): + writer(chunk) + chunks.append(chunk) + return chunks + else: + + @functools.wraps(func) + async def agen_wrapper( + *args: Any, writer: StreamWriter, **kwargs: Any + ) -> Any: + chunks = [] + async for chunk in func(*args, **kwargs): + writer(chunk) + chunks.append(chunk) + return chunks + + # Create a new parameter for the writer argument + extra_param = inspect.Parameter( + "writer", + inspect.Parameter.KEYWORD_ONLY, + # The extra argument is a keyword-only argument + default=lambda _: None, + ) + # Update the function's signature to include the extra argument + new_params = list(original_sig.parameters.values()) + [extra_param] + new_sig = original_sig.replace(parameters=new_params) + # Update the signature of the wrapper function + agen_wrapper.__signature__ = new_sig # type: ignore bound = get_runnable_for_func(agen_wrapper) stream_mode = "custom" @@ -134,6 +210,23 @@ def entrypoint( bound = get_runnable_for_func(func) stream_mode = "updates" + # get input and output types + sig = inspect.signature(func) + first_parameter_name = next(iter(sig.parameters.keys()), None) + if not first_parameter_name: + raise ValueError("Entrypoint function must have at least one parameter") + input_type = ( + sig.parameters[first_parameter_name].annotation + if sig.parameters[first_parameter_name].annotation + is not inspect.Signature.empty + else Any + ) + output_type = ( + sig.return_annotation + if sig.return_annotation is not inspect.Signature.empty + else Any + ) + return Pregel( nodes={ func.__name__: PregelNode( @@ -143,13 +236,18 @@ def entrypoint( writers=[ChannelWrite([ChannelWriteEntry(END)], tags=[TAG_HIDDEN])], ) }, - channels={START: EphemeralValue(Any), END: LastValue(Any, END)}, + channels={ + START: EphemeralValue(input_type), + END: LastValue(output_type, END), + }, input_channels=START, output_channels=END, stream_channels=END, stream_mode=stream_mode, + stream_eager=True, checkpointer=checkpointer, store=store, + config_type=config_schema, ) return _imp diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 3fbb31229..16bbf7d48 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -847,14 +847,10 @@ def _control_branch(value: Any) -> Sequence[Union[str, Send]]: commands: list[Command] = [] if isinstance(value, Command): commands.append(value) - elif ( - isinstance(value, (list, tuple)) - and value - and all(isinstance(i, Command) for i in value) - ): - commands.extend(value) - else: - return EMPTY_SEQ + elif isinstance(value, (list, tuple)): + for cmd in value: + if isinstance(cmd, Command): + commands.append(cmd) rtn: list[Union[str, Send]] = [] for command in commands: if command.graph == Command.PARENT: @@ -874,14 +870,10 @@ async def _acontrol_branch(value: Any) -> Sequence[Union[str, Send]]: commands: list[Command] = [] if isinstance(value, Command): commands.append(value) - elif ( - isinstance(value, (list, tuple)) - and value - and all(isinstance(i, Command) for i in value) - ): - commands.extend(value) - else: - return EMPTY_SEQ + elif isinstance(value, (list, tuple)): + for cmd in value: + if isinstance(cmd, Command): + commands.append(cmd) rtn: list[Union[str, Send]] = [] for command in commands: if command.graph == Command.PARENT: diff --git a/libs/langgraph/langgraph/prebuilt/interrupt.py b/libs/langgraph/langgraph/prebuilt/interrupt.py new file mode 100644 index 000000000..2ab3ee0d1 --- /dev/null +++ b/libs/langgraph/langgraph/prebuilt/interrupt.py @@ -0,0 +1,94 @@ +from typing import ( + Literal, + Optional, + Union, +) + +from typing_extensions import TypedDict + + +class HumanInterruptConfig(TypedDict): + """Configuration that defines what actions are allowed for a human interrupt. + + This controls the available interaction options when the graph is paused for human input. + + Attributes: + allow_ignore: Whether the human can choose to ignore/skip the current step + allow_respond: Whether the human can provide a text response/feedback + allow_edit: Whether the human can edit the provided content/state + allow_accept: Whether the human can accept/approve the current state + """ + + allow_ignore: bool + allow_respond: bool + allow_edit: bool + allow_accept: bool + + +class ActionRequest(TypedDict): + """Represents a request for human action within the graph execution. + + Contains the action type and any associated arguments needed for the action. + + Attributes: + action: The type or name of action being requested (e.g., "Approve XYZ action") + args: Key-value pairs of arguments needed for the action + """ + + action: str + args: dict + + +class HumanInterrupt(TypedDict): + """Represents an interrupt triggered by the graph that requires human intervention. + + This is passed to the `interrupt` function when execution is paused for human input. + + Attributes: + action_request: The specific action being requested from the human + config: Configuration defining what actions are allowed + description: Optional detailed description of what input is needed + + Example: + ```python + # Extract a tool call from the state and create an interrupt request + request = HumanInterrupt( + action_request=ActionRequest( + action="run_command", # The action being requested + args={"command": "ls", "args": ["-l"]} # Arguments for the action + ), + config=HumanInterruptConfig( + allow_ignore=True, # Allow skipping this step + allow_respond=True, # Allow text feedback + allow_edit=False, # Don't allow editing + allow_accept=True # Allow direct acceptance + ), + description="Please review the command before execution" + ) + # Send the interrupt request and get the response + response = interrupt([request])[0] + ``` + """ + + action_request: ActionRequest + config: HumanInterruptConfig + description: Optional[str] + + +class HumanResponse(TypedDict): + """The response provided by a human to an interrupt, which is returned when graph execution resumes. + + Attributes: + type: The type of response: + - "accept": Approves the current state without changes + - "ignore": Skips/ignores the current step + - "response": Provides text feedback or instructions + - "edit": Modifies the current state/content + arg: The response payload: + - None: For ignore/accept actions + - str: For text responses + - ActionRequest: For edit actions with updated content + """ + + type: Literal["accept", "ignore", "response", "edit"] + args: Union[None, str, ActionRequest] diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index cc8786968..877ddcf4e 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -115,6 +115,7 @@ from langgraph.utils.config import ( patch_checkpoint_map, patch_config, patch_configurable, + recast_checkpoint_ns, ) from langgraph.utils.fields import get_enhanced_type_hints from langgraph.utils.pydantic import create_model @@ -203,6 +204,10 @@ class Pregel(PregelProtocol): stream_mode: StreamMode = "values" """Mode to stream output, defaults to 'values'.""" + stream_eager: bool = False + """Whether to force emitting stream events eagerly, automatically turned on + for stream_mode "messages" and "custom".""" + output_channels: Union[str, Sequence[str]] stream_channels: Optional[Union[str, Sequence[str]]] = None @@ -242,6 +247,7 @@ class Pregel(PregelProtocol): channels: Optional[dict[str, Union[BaseChannel, ManagedValueSpec]]], auto_validate: bool = True, stream_mode: StreamMode = "values", + stream_eager: bool = False, output_channels: Union[str, Sequence[str]], stream_channels: Optional[Union[str, Sequence[str]]] = None, interrupt_after_nodes: Union[All, Sequence[str]] = (), @@ -259,6 +265,7 @@ class Pregel(PregelProtocol): self.nodes = nodes self.channels = channels or {} self.stream_mode = stream_mode + self.stream_eager = stream_eager self.output_channels = output_channels self.stream_channels = stream_channels self.interrupt_after_nodes = interrupt_after_nodes @@ -494,7 +501,9 @@ class Pregel(PregelProtocol): saved.metadata.get("step", -1) + 1, for_execution=True, store=self.store, - checkpointer=self.checkpointer or None, + checkpointer=self.checkpointer + if isinstance(self.checkpointer, BaseCheckpointSaver) + else None, manager=None, ) # get the subgraphs @@ -606,7 +615,9 @@ class Pregel(PregelProtocol): saved.metadata.get("step", -1) + 1, for_execution=True, store=self.store, - checkpointer=self.checkpointer or None, + checkpointer=self.checkpointer + if isinstance(self.checkpointer, BaseCheckpointSaver) + else None, manager=None, ) # get the subgraphs @@ -690,19 +701,15 @@ class Pregel(PregelProtocol): checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "") ) and CONFIG_KEY_CHECKPOINTER not in config[CONF]: # remove task_ids from checkpoint_ns - recast_checkpoint_ns = NS_SEP.join( - part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP) - ) + recast = recast_checkpoint_ns(checkpoint_ns) # find the subgraph with the matching name - for _, pregel in self.get_subgraphs( - namespace=recast_checkpoint_ns, recurse=True - ): + for _, pregel in self.get_subgraphs(namespace=recast, recurse=True): return pregel.get_state( patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}), subgraphs=subgraphs, ) else: - raise ValueError(f"Subgraph {recast_checkpoint_ns} not found") + raise ValueError(f"Subgraph {recast} not found") config = merge_configs(self.config, config) if self.config else config saved = checkpointer.get_tuple(config) @@ -727,19 +734,15 @@ class Pregel(PregelProtocol): checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "") ) and CONFIG_KEY_CHECKPOINTER not in config[CONF]: # remove task_ids from checkpoint_ns - recast_checkpoint_ns = NS_SEP.join( - part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP) - ) + recast = recast_checkpoint_ns(checkpoint_ns) # find the subgraph with the matching name - async for _, pregel in self.aget_subgraphs( - namespace=recast_checkpoint_ns, recurse=True - ): + async for _, pregel in self.aget_subgraphs(namespace=recast, recurse=True): return await pregel.aget_state( patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}), subgraphs=subgraphs, ) else: - raise ValueError(f"Subgraph {recast_checkpoint_ns} not found") + raise ValueError(f"Subgraph {recast} not found") config = merge_configs(self.config, config) if self.config else config saved = await checkpointer.aget_tuple(config) @@ -770,13 +773,9 @@ class Pregel(PregelProtocol): checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "") ) and CONFIG_KEY_CHECKPOINTER not in config[CONF]: # remove task_ids from checkpoint_ns - recast_checkpoint_ns = NS_SEP.join( - part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP) - ) + recast = recast_checkpoint_ns(checkpoint_ns) # find the subgraph with the matching name - for _, pregel in self.get_subgraphs( - namespace=recast_checkpoint_ns, recurse=True - ): + for _, pregel in self.get_subgraphs(namespace=recast, recurse=True): yield from pregel.get_state_history( patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}), filter=filter, @@ -785,7 +784,7 @@ class Pregel(PregelProtocol): ) return else: - raise ValueError(f"Subgraph {recast_checkpoint_ns} not found") + raise ValueError(f"Subgraph {recast} not found") config = merge_configs( self.config, @@ -820,13 +819,9 @@ class Pregel(PregelProtocol): checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "") ) and CONFIG_KEY_CHECKPOINTER not in config[CONF]: # remove task_ids from checkpoint_ns - recast_checkpoint_ns = NS_SEP.join( - part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP) - ) + recast = recast_checkpoint_ns(checkpoint_ns) # find the subgraph with the matching name - async for _, pregel in self.aget_subgraphs( - namespace=recast_checkpoint_ns, recurse=True - ): + async for _, pregel in self.aget_subgraphs(namespace=recast, recurse=True): async for state in pregel.aget_state_history( patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}), filter=filter, @@ -836,7 +831,7 @@ class Pregel(PregelProtocol): yield state return else: - raise ValueError(f"Subgraph {recast_checkpoint_ns} not found") + raise ValueError(f"Subgraph {recast} not found") config = merge_configs( self.config, @@ -875,20 +870,16 @@ class Pregel(PregelProtocol): checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "") ) and CONFIG_KEY_CHECKPOINTER not in config[CONF]: # remove task_ids from checkpoint_ns - recast_checkpoint_ns = NS_SEP.join( - part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP) - ) + recast = recast_checkpoint_ns(checkpoint_ns) # find the subgraph with the matching name - for _, pregel in self.get_subgraphs( - namespace=recast_checkpoint_ns, recurse=True - ): + for _, pregel in self.get_subgraphs(namespace=recast, recurse=True): return pregel.update_state( patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}), values, as_node, ) else: - raise ValueError(f"Subgraph {recast_checkpoint_ns} not found") + raise ValueError(f"Subgraph {recast} not found") # get last checkpoint config = ensure_config(self.config, config) @@ -926,7 +917,9 @@ class Pregel(PregelProtocol): saved.metadata.get("step", -1) + 1, for_execution=True, store=self.store, - checkpointer=self.checkpointer or None, + checkpointer=self.checkpointer + if isinstance(self.checkpointer, BaseCheckpointSaver) + else None, manager=None, ) # apply null writes @@ -1020,7 +1013,9 @@ class Pregel(PregelProtocol): saved.metadata.get("step", -1) + 1, for_execution=True, store=self.store, - checkpointer=self.checkpointer or None, + checkpointer=self.checkpointer + if isinstance(self.checkpointer, BaseCheckpointSaver) + else None, manager=None, ) # apply null writes @@ -1155,20 +1150,16 @@ class Pregel(PregelProtocol): checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "") ) and CONFIG_KEY_CHECKPOINTER not in config[CONF]: # remove task_ids from checkpoint_ns - recast_checkpoint_ns = NS_SEP.join( - part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP) - ) + recast = recast_checkpoint_ns(checkpoint_ns) # find the subgraph with the matching name - async for _, pregel in self.aget_subgraphs( - namespace=recast_checkpoint_ns, recurse=True - ): + async for _, pregel in self.aget_subgraphs(namespace=recast, recurse=True): return await pregel.aupdate_state( patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}), values, as_node, ) else: - raise ValueError(f"Subgraph {recast_checkpoint_ns} not found") + raise ValueError(f"Subgraph {recast} not found") # get last checkpoint config = ensure_config(self.config, config) @@ -1209,7 +1200,9 @@ class Pregel(PregelProtocol): saved.metadata.get("step", -1) + 1, for_execution=True, store=self.store, - checkpointer=self.checkpointer or None, + checkpointer=self.checkpointer + if isinstance(self.checkpointer, BaseCheckpointSaver) + else None, manager=None, ) # apply null writes @@ -1303,7 +1296,9 @@ class Pregel(PregelProtocol): saved.metadata.get("step", -1) + 1, for_execution=True, store=self.store, - checkpointer=self.checkpointer or None, + checkpointer=self.checkpointer + if isinstance(self.checkpointer, BaseCheckpointSaver) + else None, manager=None, ) # apply null writes @@ -1455,6 +1450,8 @@ class Pregel(PregelProtocol): checkpointer: Optional[BaseCheckpointSaver] = None elif CONFIG_KEY_CHECKPOINTER in config.get(CONF, {}): checkpointer = config[CONF][CONFIG_KEY_CHECKPOINTER] + elif self.checkpointer is True: + raise RuntimeError("checkpointer=True cannot be used for root graphs.") else: checkpointer = self.checkpointer if checkpointer and not config.get(CONF): @@ -1598,6 +1595,12 @@ class Pregel(PregelProtocol): interrupt_after=interrupt_after, debug=debug, ) + # set up subgraph checkpointing + if self.checkpointer is True: + ns = cast(str, config[CONF][CONFIG_KEY_CHECKPOINT_NS]) + config[CONF][CONFIG_KEY_CHECKPOINT_NS] = NS_SEP.join( + part.split(NS_END)[0] for part in ns.split(NS_SEP) + ) # set up messages stream mode if "messages" in stream_modes: run_manager.inheritable_handlers.append( @@ -1634,7 +1637,12 @@ class Pregel(PregelProtocol): if subgraphs: loop.config[CONF][CONFIG_KEY_STREAM] = loop.stream # enable concurrent streaming - if subgraphs or "messages" in stream_modes or "custom" in stream_modes: + if ( + self.stream_eager + or subgraphs + or "messages" in stream_modes + or "custom" in stream_modes + ): # we are careful to have a single waiter live at any one time # because on exit we increment semaphore count by exactly 1 waiter: Optional[concurrent.futures.Future] = None @@ -1864,7 +1872,12 @@ class Pregel(PregelProtocol): stream_put, stream_modes ) # enable concurrent streaming - if subgraphs or "messages" in stream_modes or "custom" in stream_modes: + if ( + self.stream_eager + or subgraphs + or "messages" in stream_modes + or "custom" in stream_modes + ): def get_waiter() -> asyncio.Task[None]: return aioloop.create_task(stream.wait()) diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 12a53b8ad..d1247dbc4 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -37,15 +37,16 @@ from langgraph.constants import ( CONFIG_KEY_CHECKPOINT_MAP, CONFIG_KEY_CHECKPOINT_NS, CONFIG_KEY_CHECKPOINTER, + CONFIG_KEY_END, CONFIG_KEY_READ, CONFIG_KEY_SCRATCHPAD, CONFIG_KEY_SEND, CONFIG_KEY_STORE, CONFIG_KEY_TASK_ID, - CONFIG_KEY_WRITES, EMPTY_SEQ, ERROR, INTERRUPT, + MISSING, NO_WRITES, NS_END, NS_SEP, @@ -71,6 +72,7 @@ from langgraph.types import ( All, LoopProtocol, PregelExecutableTask, + PregelScratchpad, PregelTask, RetryPolicy, ) @@ -502,13 +504,13 @@ def prepare_single_task( }, CONFIG_KEY_CHECKPOINT_ID: None, CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns, - CONFIG_KEY_WRITES: [ - w - for w in pending_writes - + configurable.get(CONFIG_KEY_WRITES, []) - if w[0] in (NULL_TASK_ID, task_id) - ], - CONFIG_KEY_SCRATCHPAD: {}, + CONFIG_KEY_SCRATCHPAD: _scratchpad( + pending_writes, + task_id, + ), + CONFIG_KEY_END: checkpoint["channel_values"].get( + "__end__", None + ), }, ), triggers, @@ -614,13 +616,13 @@ def prepare_single_task( }, CONFIG_KEY_CHECKPOINT_ID: None, CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns, - CONFIG_KEY_WRITES: [ - w - for w in pending_writes - + configurable.get(CONFIG_KEY_WRITES, []) - if w[0] in (NULL_TASK_ID, task_id) - ], - CONFIG_KEY_SCRATCHPAD: {}, + CONFIG_KEY_SCRATCHPAD: _scratchpad( + pending_writes, + task_id, + ), + CONFIG_KEY_END: checkpoint["channel_values"].get( + "__end__", None + ), }, ), triggers, @@ -685,7 +687,7 @@ def prepare_single_task( "langgraph_checkpoint_ns": task_checkpoint_ns, } if task_id_checksum is not None: - assert task_id == task_id_checksum + assert task_id == task_id_checksum, f"{task_id} != {task_id_checksum}" if for_execution: if node := proc.node: if proc.metadata: @@ -738,13 +740,13 @@ def prepare_single_task( }, CONFIG_KEY_CHECKPOINT_ID: None, CONFIG_KEY_CHECKPOINT_NS: task_checkpoint_ns, - CONFIG_KEY_WRITES: [ - w - for w in pending_writes - + configurable.get(CONFIG_KEY_WRITES, []) - if w[0] in (NULL_TASK_ID, task_id) - ], - CONFIG_KEY_SCRATCHPAD: {}, + CONFIG_KEY_SCRATCHPAD: _scratchpad( + pending_writes, + task_id, + ), + CONFIG_KEY_END: checkpoint["channel_values"].get( + "__end__", None + ), }, ), triggers, @@ -758,6 +760,27 @@ def prepare_single_task( return PregelTask(task_id, name, task_path[:3]) +def _scratchpad( + pending_writes: Sequence[PendingWrite], + task_id: str, +) -> PregelScratchpad: + return PregelScratchpad( + # call + call_counter=0, + # interrupt + interrupt_counter=-1, + resume=next( + (w[2] for w in pending_writes if w[0] == task_id and w[1] == RESUME), [] + ), + null_resume=next( + (w[2] for w in pending_writes if w[0] == NULL_TASK_ID and w[1] == RESUME), + MISSING, + ), + # subgraph + subgraph_counter=0, + ) + + def _proc_input( proc: PregelNode, managed: ManagedValueMapping, diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 5fd9b5d96..e92f0c267 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -47,12 +47,14 @@ from langgraph.constants import ( CONFIG_KEY_DELEGATE, CONFIG_KEY_ENSURE_LATEST, CONFIG_KEY_RESUMING, + CONFIG_KEY_SCRATCHPAD, CONFIG_KEY_STREAM, CONFIG_KEY_TASK_ID, EMPTY_SEQ, ERROR, INPUT, INTERRUPT, + MISSING, NS_SEP, NULL_TASK_ID, PUSH, @@ -61,12 +63,10 @@ from langgraph.constants import ( TAG_HIDDEN, ) from langgraph.errors import ( - _SEEN_CHECKPOINT_NS, CheckpointNotLatest, EmptyInputError, GraphDelegate, GraphInterrupt, - MultipleSubgraphsError, ) from langgraph.managed.base import ( ManagedValueMapping, @@ -114,6 +114,7 @@ from langgraph.types import ( Command, LoopProtocol, PregelExecutableTask, + PregelScratchpad, StreamChunk, StreamProtocol, ) @@ -201,7 +202,6 @@ class PregelLoop(LoopProtocol): interrupt_after: Union[All, Sequence[str]] = EMPTY_SEQ, interrupt_before: Union[All, Sequence[str]] = EMPTY_SEQ, manager: Union[None, AsyncParentRunManager, ParentRunManager] = None, - check_subgraphs: bool = True, debug: bool = False, ) -> None: super().__init__( @@ -228,20 +228,26 @@ class PregelLoop(LoopProtocol): self.debug = debug if self.stream is not None and CONFIG_KEY_STREAM in config[CONF]: self.stream = DuplexStream(self.stream, config[CONF][CONFIG_KEY_STREAM]) + scratchpad: Optional[PregelScratchpad] = config[CONF].get(CONFIG_KEY_SCRATCHPAD) + if not self.config[CONF].get(CONFIG_KEY_DELEGATE) and scratchpad is not None: + if scratchpad["subgraph_counter"]: + self.config = patch_configurable( + self.config, + { + CONFIG_KEY_CHECKPOINT_NS: NS_SEP.join( + ( + config[CONF][CONFIG_KEY_CHECKPOINT_NS], + str(scratchpad["subgraph_counter"]), + ) + ) + }, + ) + scratchpad["subgraph_counter"] += 1 if not self.is_nested and config[CONF].get(CONFIG_KEY_CHECKPOINT_NS): self.config = patch_configurable( self.config, {CONFIG_KEY_CHECKPOINT_NS: "", CONFIG_KEY_CHECKPOINT_ID: None}, ) - if check_subgraphs and self.is_nested and self.checkpointer is not None: - if self.config[CONF][CONFIG_KEY_CHECKPOINT_NS] in _SEEN_CHECKPOINT_NS: - raise MultipleSubgraphsError( - "Multiple subgraphs called inside the same node\n\n" - "Troubleshooting URL: https://python.langchain.com/docs" - "/troubleshooting/errors/MULTIPLE_SUBGRAPHS/" - ) - else: - _SEEN_CHECKPOINT_NS.add(self.config[CONF][CONFIG_KEY_CHECKPOINT_NS]) if ( CONFIG_KEY_CHECKPOINT_MAP in self.config[CONF] and self.config[CONF].get(CONFIG_KEY_CHECKPOINT_NS) @@ -339,11 +345,11 @@ class PregelLoop(LoopProtocol): (PUSH, task.path, write_idx, task.id, call), None, checkpoint=self.checkpoint, - pending_writes=[(task.id, *w) for w in task.writes], + pending_writes=self.checkpoint_pending_writes, processes=self.nodes, channels=self.channels, managed=self.managed, - config=self.config, + config=task.config, step=self.step, for_execution=True, store=self.store, @@ -556,8 +562,16 @@ class PregelLoop(LoopProtocol): ) ) + # take resume value from parent + if scratchpad := configurable.get(CONFIG_KEY_SCRATCHPAD): + if scratchpad["null_resume"] is not MISSING: + self.put_writes(NULL_TASK_ID, [(RESUME, scratchpad["null_resume"])]) # map command to writes if isinstance(self.input, Command): + if self.input.resume is not None and not self.checkpointer: + raise RuntimeError( + "Cannot use Command(resume=...) without checkpointer" + ) writes: defaultdict[str, list[tuple[str, Any]]] = defaultdict(list) # group writes by task ID for tid, c, v in map_command(self.input, self.checkpoint_pending_writes): @@ -807,7 +821,6 @@ class SyncPregelLoop(PregelLoop, ContextManager): interrupt_before: Union[All, Sequence[str]] = EMPTY_SEQ, output_keys: Union[str, Sequence[str]] = EMPTY_SEQ, stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ, - check_subgraphs: bool = True, debug: bool = False, ) -> None: super().__init__( @@ -822,7 +835,6 @@ class SyncPregelLoop(PregelLoop, ContextManager): stream_keys=stream_keys, interrupt_after=interrupt_after, interrupt_before=interrupt_before, - check_subgraphs=check_subgraphs, manager=manager, debug=debug, ) @@ -944,7 +956,6 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): manager: Union[None, AsyncParentRunManager, ParentRunManager] = None, output_keys: Union[str, Sequence[str]] = EMPTY_SEQ, stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ, - check_subgraphs: bool = True, debug: bool = False, ) -> None: super().__init__( @@ -959,7 +970,6 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): stream_keys=stream_keys, interrupt_after=interrupt_after, interrupt_before=interrupt_before, - check_subgraphs=check_subgraphs, manager=manager, debug=debug, ) diff --git a/libs/langgraph/langgraph/pregel/retry.py b/libs/langgraph/langgraph/pregel/retry.py index 29faaab21..6d0e43b54 100644 --- a/libs/langgraph/langgraph/pregel/retry.py +++ b/libs/langgraph/langgraph/pregel/retry.py @@ -12,7 +12,7 @@ from langgraph.constants import ( CONFIG_KEY_RESUMING, NS_SEP, ) -from langgraph.errors import _SEEN_CHECKPOINT_NS, GraphBubbleUp, ParentCommand +from langgraph.errors import GraphBubbleUp, ParentCommand from langgraph.types import Command, PregelExecutableTask, RetryPolicy from langgraph.utils.config import patch_configurable @@ -48,7 +48,10 @@ def run_with_retry( break elif cmd.graph == Command.PARENT: # this command is for the parent graph, assign it to the parent - parent_ns = NS_SEP.join(ns.split(NS_SEP)[:-1]) + parts = ns.split(NS_SEP) + if parts[-1].isdigit(): + parts.pop() + parent_ns = NS_SEP.join(parts[:-1]) exc.args = (replace(cmd, graph=parent_ns),) # bubble up raise @@ -96,13 +99,6 @@ def run_with_retry( ) # signal subgraphs to resume (if available) config = patch_configurable(config, {CONFIG_KEY_RESUMING: True}) - # clear checkpoint_ns seen (for subgraph detection) - if checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS): - _SEEN_CHECKPOINT_NS.discard(checkpoint_ns) - finally: - # clear checkpoint_ns seen (for subgraph detection) - if checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS): - _SEEN_CHECKPOINT_NS.discard(checkpoint_ns) async def arun_with_retry( @@ -140,7 +136,10 @@ async def arun_with_retry( break elif cmd.graph == Command.PARENT: # this command is for the parent graph, assign it to the parent - parent_ns = NS_SEP.join(ns.split(NS_SEP)[:-1]) + parts = ns.split(NS_SEP) + if parts[-1].isdigit(): + parts.pop() + parent_ns = NS_SEP.join(parts[:-1]) exc.args = (replace(cmd, graph=parent_ns),) # bubble up raise @@ -188,10 +187,3 @@ async def arun_with_retry( ) # signal subgraphs to resume (if available) config = patch_configurable(config, {CONFIG_KEY_RESUMING: True}) - # clear checkpoint_ns seen (for subgraph detection) - if checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS): - _SEEN_CHECKPOINT_NS.discard(checkpoint_ns) - finally: - # clear checkpoint_ns seen (for subgraph detection) - if checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS): - _SEEN_CHECKPOINT_NS.discard(checkpoint_ns) diff --git a/libs/langgraph/langgraph/pregel/runner.py b/libs/langgraph/langgraph/pregel/runner.py index 790d22a06..d354c7866 100644 --- a/libs/langgraph/langgraph/pregel/runner.py +++ b/libs/langgraph/langgraph/pregel/runner.py @@ -107,12 +107,11 @@ class PregelRunner: elif next_task.writes: # if it already ran, return the result fut = concurrent.futures.Future() - if ( - val := next( - (v for c, v in next_task.writes if c == RETURN), MISSING - ) - ) and val is not MISSING: - fut.set_result(val) + ret = next( + (v for c, v in next_task.writes if c == RETURN), MISSING + ) + if ret is not MISSING: + fut.set_result(ret) elif exc := next( (v for c, v in next_task.writes if c == ERROR), None ): @@ -295,12 +294,11 @@ class PregelRunner: elif next_task.writes: # if it already ran, return the result fut = asyncio.Future() - if ( - val := next( - (v for c, v in next_task.writes if c == RETURN), MISSING - ) - ) and val is not MISSING: - fut.set_result(val) + ret = next( + (v for c, v in next_task.writes if c == RETURN), MISSING + ) + if ret is not MISSING: + fut.set_result(ret) elif exc := next( (v for c, v in next_task.writes if c == ERROR), None ): diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 0b9fb9b1b..09f777d01 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -21,11 +21,7 @@ from typing import ( from langchain_core.runnables import Runnable, RunnableConfig from typing_extensions import Self, TypedDict -from langgraph.checkpoint.base import ( - BaseCheckpointSaver, - CheckpointMetadata, - PendingWrite, -) +from langgraph.checkpoint.base import BaseCheckpointSaver, CheckpointMetadata if TYPE_CHECKING: from langgraph.store.base import BaseStore @@ -42,9 +38,11 @@ except ImportError: All = Literal["*"] """Special value to indicate that graph should interrupt on all nodes.""" -Checkpointer = Union[None, Literal[False], BaseCheckpointSaver] -"""Type of the checkpointer to use for a subgraph. False disables checkpointing, -even if the parent graph has a checkpointer. None inherits checkpointer.""" +Checkpointer = Union[None, bool, BaseCheckpointSaver] +"""Type of the checkpointer to use for a subgraph. +- True enables persistent checkpointing for this subgraph. +- False disables checkpointing, even if the parent graph has a checkpointer. +- None inherits checkpointer from the parent graph.""" StreamMode = Literal["values", "updates", "debug", "messages", "custom"] """How the stream method should emit outputs. @@ -341,13 +339,15 @@ class LoopProtocol: self.stop = stop -class PregelScratchpad(TypedDict, total=False): - # interrupt - interrupt_counter: int - used_null_resume: bool - resume: list[Any] +class PregelScratchpad(TypedDict): # call call_counter: int + # interrupt + interrupt_counter: int + resume: list[Any] + null_resume: Any + # subgraph + subgraph_counter: int def interrupt(value: Any) -> Any: @@ -449,10 +449,8 @@ def interrupt(value: Any) -> Any: CONFIG_KEY_CHECKPOINT_NS, CONFIG_KEY_SCRATCHPAD, CONFIG_KEY_SEND, - CONFIG_KEY_TASK_ID, - CONFIG_KEY_WRITES, + MISSING, NS_SEP, - NULL_TASK_ID, RESUME, ) from langgraph.errors import GraphInterrupt @@ -461,29 +459,20 @@ def interrupt(value: Any) -> Any: conf = get_config()["configurable"] # track interrupt index scratchpad: PregelScratchpad = conf[CONFIG_KEY_SCRATCHPAD] - if "interrupt_counter" not in scratchpad: - scratchpad["interrupt_counter"] = 0 - else: - scratchpad["interrupt_counter"] += 1 + scratchpad["interrupt_counter"] += 1 idx = scratchpad["interrupt_counter"] # find previous resume values - task_id = conf[CONFIG_KEY_TASK_ID] - writes: list[PendingWrite] = conf[CONFIG_KEY_WRITES] - scratchpad.setdefault( - "resume", next((w[2] for w in writes if w[0] == task_id and w[1] == RESUME), []) - ) if scratchpad["resume"]: if idx < len(scratchpad["resume"]): return scratchpad["resume"][idx] # find current resume value - if not scratchpad.get("used_null_resume"): - scratchpad["used_null_resume"] = True - for tid, c, v in sorted(writes, key=lambda x: x[0], reverse=True): - if tid == NULL_TASK_ID and c == RESUME: - assert len(scratchpad["resume"]) == idx, (scratchpad["resume"], idx) - scratchpad["resume"].append(v) - conf[CONFIG_KEY_SEND]([(RESUME, scratchpad["resume"])]) - return v + if scratchpad["null_resume"] is not MISSING: + assert len(scratchpad["resume"]) == idx, (scratchpad["resume"], idx) + v = scratchpad["null_resume"] + scratchpad["null_resume"] = MISSING + scratchpad["resume"].append(v) + conf[CONFIG_KEY_SEND]([(RESUME, scratchpad["resume"])]) + return v # no resume value found raise GraphInterrupt( ( diff --git a/libs/langgraph/langgraph/utils/config.py b/libs/langgraph/langgraph/utils/config.py index ac803cc35..309c6d6be 100644 --- a/libs/langgraph/langgraph/utils/config.py +++ b/libs/langgraph/langgraph/utils/config.py @@ -23,9 +23,25 @@ from langgraph.constants import ( CONFIG_KEY_CHECKPOINT_ID, CONFIG_KEY_CHECKPOINT_MAP, CONFIG_KEY_CHECKPOINT_NS, + NS_END, + NS_SEP, ) +def recast_checkpoint_ns(ns: str) -> str: + """Remove task IDs from checkpoint namespace. + + Args: + ns (str): The checkpoint namespace with task IDs. + + Returns: + str: The checkpoint namespace without task IDs. + """ + return NS_SEP.join( + part.split(NS_END)[0] for part in ns.split(NS_SEP) if not part.isdigit() + ) + + def patch_configurable( config: Optional[RunnableConfig], patch: dict[str, Any] ) -> RunnableConfig: diff --git a/libs/langgraph/langgraph/utils/runnable.py b/libs/langgraph/langgraph/utils/runnable.py index 9f1a7ef8d..19a70e4cf 100644 --- a/libs/langgraph/langgraph/utils/runnable.py +++ b/libs/langgraph/langgraph/utils/runnable.py @@ -34,7 +34,12 @@ from langchain_core.runnables.utils import Input from langchain_core.tracers._streaming import _StreamingCallbackHandler from typing_extensions import TypeGuard -from langgraph.constants import CONF, CONFIG_KEY_STORE, CONFIG_KEY_STREAM_WRITER +from langgraph.constants import ( + CONF, + CONFIG_KEY_END, + CONFIG_KEY_STORE, + CONFIG_KEY_STREAM_WRITER, +) from langgraph.store.base import BaseStore from langgraph.types import StreamWriter from langgraph.utils.config import ( @@ -58,6 +63,10 @@ class StrEnum(str, enum.Enum): """A string enum.""" +# Special type to denote any type is accepted +ANY_TYPE = object() + + ASYNCIO_ACCEPTS_CONTEXT = sys.version_info >= (3, 11) KWARGS_CONFIG_KEYS: tuple[tuple[str, tuple[Any, ...], str, Any], ...] = ( @@ -73,6 +82,12 @@ KWARGS_CONFIG_KEYS: tuple[tuple[str, tuple[Any, ...], str, Any], ...] = ( CONFIG_KEY_STORE, inspect.Parameter.empty, ), + ( + sys.intern("previous"), + (ANY_TYPE,), + CONFIG_KEY_END, + inspect.Parameter.empty, + ), ) """List of kwargs that can be passed to functions, and their corresponding config keys, default values and type annotations. @@ -135,9 +150,12 @@ class RunnableCallable(Runnable): self.func_accepts: dict[str, bool] = {} for kw, typ, _, _ in KWARGS_CONFIG_KEYS: p = params.get(kw) - self.func_accepts[kw] = ( - p is not None and p.annotation in typ and p.kind in VALID_KINDS - ) + if typ == (ANY_TYPE,): + self.func_accepts[kw] = p is not None and p.kind in VALID_KINDS + else: + self.func_accepts[kw] = ( + p is not None and p.annotation in typ and p.kind in VALID_KINDS + ) def __repr__(self) -> str: repr_args = { @@ -162,16 +180,20 @@ class RunnableCallable(Runnable): if self.func_accepts_config: kwargs["config"] = config _conf = config[CONF] - for kw, _, ck, defv in KWARGS_CONFIG_KEYS: + for kw, _, config_key, default_value in KWARGS_CONFIG_KEYS: if not self.func_accepts[kw]: continue - if defv is inspect.Parameter.empty and kw not in kwargs and ck not in _conf: + if ( + default_value is inspect.Parameter.empty + and kw not in kwargs + and config_key not in _conf + ): raise ValueError( - f"Missing required config key '{ck}' for '{self.name}'." + f"Missing required config key '{config_key}' for '{self.name}'." ) elif kwargs.get(kw) is None: - kwargs[kw] = _conf.get(ck, defv) + kwargs[kw] = _conf.get(config_key, default_value) context = copy_context() if self.trace: @@ -210,16 +232,20 @@ class RunnableCallable(Runnable): if self.func_accepts_config: kwargs["config"] = config _conf = config[CONF] - for kw, _, ck, defv in KWARGS_CONFIG_KEYS: + for kw, _, config_key, default_value in KWARGS_CONFIG_KEYS: if not self.func_accepts[kw]: continue - if defv is inspect.Parameter.empty and kw not in kwargs and ck not in _conf: + if ( + default_value is inspect.Parameter.empty + and kw not in kwargs + and config_key not in _conf + ): raise ValueError( - f"Missing required config key '{ck}' for '{self.name}'." + f"Missing required config key '{config_key}' for '{self.name}'." ) elif kwargs.get(kw) is None: - kwargs[kw] = _conf.get(ck, defv) + kwargs[kw] = _conf.get(config_key, default_value) context = copy_context() if self.trace: callback_manager = get_async_callback_manager_for_config(config, self.tags) diff --git a/libs/langgraph/pyproject.toml b/libs/langgraph/pyproject.toml index 04c338c56..c206bcabf 100644 --- a/libs/langgraph/pyproject.toml +++ b/libs/langgraph/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph" -version = "0.2.62" +version = "0.2.63" description = "Building stateful, multi-actor applications with LLMs" authors = [] license = "MIT" diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 12b7c6863..a30a66cba 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -51,7 +51,7 @@ from langgraph.checkpoint.base import ( ) from langgraph.checkpoint.memory import MemorySaver from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, START -from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError +from langgraph.errors import InvalidUpdateError from langgraph.func import entrypoint, task from langgraph.graph import END, Graph, StateGraph from langgraph.graph.message import MessageGraph, MessagesState, add_messages @@ -1436,6 +1436,9 @@ def test_imp_task(request: pytest.FixtureRequest, checkpointer_name: str) -> Non checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") mapper_calls = 0 + class Config: + model: str + @task() def mapper(input: int) -> str: nonlocal mapper_calls @@ -1443,13 +1446,57 @@ def test_imp_task(request: pytest.FixtureRequest, checkpointer_name: str) -> Non time.sleep(input / 100) return str(input) * 2 - @entrypoint(checkpointer=checkpointer) + @entrypoint(checkpointer=checkpointer, config_schema=Config) def graph(input: list[int]) -> list[str]: futures = [mapper(i) for i in input] mapped = [f.result() for f in futures] answer = interrupt("question") return [m + answer for m in mapped] + assert graph.get_input_jsonschema() == { + "type": "array", + "items": {"type": "integer"}, + "title": "LangGraphInput", + } + assert graph.get_output_jsonschema() == { + "type": "array", + "items": {"type": "string"}, + "title": "LangGraphOutput", + } + assert graph.get_config_jsonschema() == { + "$defs": { + "Configurable": { + "properties": { + "model": {"default": None, "title": "Model", "type": "string"}, + "checkpoint_id": { + "anyOf": [{"type": "string"}, {"type": "null"}], + "default": None, + "description": "Pass to fetch a past checkpoint. If None, fetches the latest checkpoint.", + "title": "Checkpoint ID", + }, + "checkpoint_ns": { + "default": "", + "description": 'Checkpoint namespace. Denotes the path to the subgraph node the checkpoint originates from, separated by `|` character, e.g. `"child|grandchild"`. Defaults to "" (root graph).', + "title": "Checkpoint NS", + "type": "string", + }, + "thread_id": { + "default": "", + "title": "Thread ID", + "type": "string", + }, + }, + "title": "Configurable", + "type": "object", + } + }, + "properties": { + "configurable": {"$ref": "#/$defs/Configurable", "default": None} + }, + "title": "LangGraphConfig", + "type": "object", + } + thread1 = {"configurable": {"thread_id": "1"}} assert [*graph.stream([0, 1], thread1)] == [ {"mapper": "00"}, @@ -1745,9 +1792,8 @@ def test_invoke_join_then_call_other_pregel( # add checkpointer app.checkpointer = checkpointer - # subgraph is called twice in the same node, through .map(), so raises - with pytest.raises(MultipleSubgraphsError): - app.invoke([2, 3], {"configurable": {"thread_id": "1"}}) + # subgraph is called twice in the same node, but that works + assert app.invoke([2, 3], {"configurable": {"thread_id": "1"}}) == 27 # set inner graph checkpointer NeverCheckpoint inner_app.checkpointer = False @@ -2167,10 +2213,10 @@ def test_in_one_fan_out_state_graph_waiting_edge( @workflow.add_node def rewrite_query(data: State) -> State: - return {"query": f'query: {data["query"]}'} + return {"query": f"query: {data['query']}"} def analyzer_one(data: State) -> State: - return {"query": f'analyzed: {data["query"]}'} + return {"query": f"analyzed: {data['query']}"} def retriever_one(data: State) -> State: return {"docs": ["doc1", "doc2"]} @@ -2307,10 +2353,10 @@ def test_in_one_fan_out_state_graph_waiting_edge_via_branch( docs: Annotated[list[str], sorted_add] def rewrite_query(data: State) -> State: - return {"query": f'query: {data["query"]}'} + return {"query": f"query: {data['query']}"} def analyzer_one(data: State) -> State: - return {"query": f'analyzed: {data["query"]}'} + return {"query": f"analyzed: {data['query']}"} def retriever_one(data: State) -> State: return {"docs": ["doc1", "doc2"]} @@ -2740,11 +2786,11 @@ def test_in_one_fan_out_state_graph_waiting_edge_plus_regular( docs: Annotated[list[str], sorted_add] def rewrite_query(data: State) -> State: - return {"query": f'query: {data["query"]}'} + return {"query": f"query: {data['query']}"} def analyzer_one(data: State) -> State: time.sleep(0.1) - return {"query": f'analyzed: {data["query"]}'} + return {"query": f"analyzed: {data['query']}"} def retriever_one(data: State) -> State: return {"docs": ["doc1", "doc2"]} @@ -2830,10 +2876,10 @@ def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None: docs: Annotated[list[str], sorted_add] def rewrite_query(data: State) -> State: - return {"query": f'query: {data["query"]}'} + return {"query": f"query: {data['query']}"} def analyzer_one(data: State) -> State: - return {"query": f'analyzed: {data["query"]}'} + return {"query": f"analyzed: {data['query']}"} def retriever_one(data: State) -> State: return {"docs": ["doc1", "doc2"]} @@ -2903,10 +2949,10 @@ def test_callable_in_conditional_edges_with_no_path_map() -> None: query: str def rewrite(data: State) -> State: - return {"query": f'query: {data["query"]}'} + return {"query": f"query: {data['query']}"} def analyze(data: State) -> State: - return {"query": f'analyzed: {data["query"]}'} + return {"query": f"analyzed: {data['query']}"} class ChooseAnalyzer: def __call__(self, data: State) -> str: @@ -2929,10 +2975,10 @@ def test_function_in_conditional_edges_with_no_path_map() -> None: query: str def rewrite(data: State) -> State: - return {"query": f'query: {data["query"]}'} + return {"query": f"query: {data['query']}"} def analyze(data: State) -> State: - return {"query": f'analyzed: {data["query"]}'} + return {"query": f"analyzed: {data['query']}"} def choose_analyzer(data: State) -> str: return "analyzer" @@ -2965,13 +3011,13 @@ def test_in_one_fan_out_state_graph_waiting_edge_multiple_cond_edge() -> None: docs: Annotated[list[str], sorted_add] def rewrite_query(data: State) -> State: - return {"query": f'query: {data["query"]}'} + return {"query": f"query: {data['query']}"} def retriever_picker(data: State) -> list[str]: return ["analyzer_one", "retriever_two"] def analyzer_one(data: State) -> State: - return {"query": f'analyzed: {data["query"]}'} + return {"query": f"analyzed: {data['query']}"} def retriever_one(data: State) -> State: return {"docs": ["doc1", "doc2"]} @@ -3189,6 +3235,66 @@ def test_nested_graph(snapshot: SnapshotAssertion) -> None: ] +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_subgraph_checkpoint_true( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) + + class InnerState(TypedDict): + my_key: Annotated[str, operator.add] + my_other_key: str + + def inner_1(state: InnerState): + return {"my_key": " got here", "my_other_key": state["my_key"]} + + def inner_2(state: InnerState): + return {"my_key": " and there"} + + inner = StateGraph(InnerState) + inner.add_node("inner_1", inner_1) + inner.add_node("inner_2", inner_2) + inner.add_edge("inner_1", "inner_2") + inner.set_entry_point("inner_1") + inner.set_finish_point("inner_2") + + class State(TypedDict): + my_key: str + + graph = StateGraph(State) + graph.add_node("inner", inner.compile(checkpointer=True)) + graph.add_edge(START, "inner") + graph.add_conditional_edges( + "inner", lambda s: "inner" if s["my_key"].count("there") < 2 else END + ) + app = graph.compile(checkpointer=checkpointer) + + config = {"configurable": {"thread_id": "2"}} + assert [c for c in app.stream({"my_key": ""}, config, subgraphs=True)] == [ + (("inner",), {"inner_1": {"my_key": " got here", "my_other_key": ""}}), + (("inner",), {"inner_2": {"my_key": " and there"}}), + ((), {"inner": {"my_key": " got here and there"}}), + ( + ("inner",), + { + "inner_1": { + "my_key": " got here", + "my_other_key": " got here and there got here and there", + } + }, + ), + (("inner",), {"inner_2": {"my_key": " and there"}}), + ( + (), + { + "inner": { + "my_key": " got here and there got here and there got here and there" + } + }, + ), + ] + + @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_stream_subgraphs_during_execution( request: pytest.FixtureRequest, checkpointer_name: str @@ -4843,6 +4949,69 @@ def test_interrupt_loop(request: pytest.FixtureRequest, checkpointer_name: str): ] +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_interrupt_functional( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer: BaseCheckpointSaver = request.getfixturevalue( + f"checkpointer_{checkpointer_name}" + ) + + @task + def foo(state: dict) -> dict: + return {"a": state["a"] + "foo"} + + @task + def bar(state: dict) -> dict: + return {"a": state["a"] + "bar", "b": state["b"]} + + @entrypoint(checkpointer=checkpointer) + def graph(inputs: dict) -> dict: + fut_foo = foo(inputs) + value = interrupt("Provide value for bar:") + bar_input = {**fut_foo.result(), "b": value} + fut_bar = bar(bar_input) + return fut_bar.result() + + config = {"configurable": {"thread_id": "1"}} + # First run, interrupted at bar + graph.invoke({"a": ""}, config) + # Resume with an answer + res = graph.invoke(Command(resume="bar"), config) + assert res == {"a": "foobar", "b": "bar"} + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_interrupt_task_functional( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer: BaseCheckpointSaver = request.getfixturevalue( + f"checkpointer_{checkpointer_name}" + ) + + @task + def foo(state: dict) -> dict: + return {"a": state["a"] + "foo"} + + @task + def bar(state: dict) -> dict: + value = interrupt("Provide value for bar:") + return {"a": state["a"] + value} + + @entrypoint(checkpointer=checkpointer) + def graph(inputs: dict) -> dict: + fut_foo = foo(inputs) + fut_bar = bar(fut_foo.result()) + return fut_bar.result() + + config = {"configurable": {"thread_id": "1"}} + # First run, interrupted at bar + graph.invoke({"a": ""}, config) + # Resume with an answer + res = graph.invoke(Command(resume="bar"), config) + assert res == {"a": "foobar"} + + def test_root_mixed_return() -> None: def my_node(state: list[str]): return [Command(update=["a"]), ["b"]] @@ -5262,9 +5431,10 @@ def test_multiple_updates() -> None: ] -def test_falsy_return_from_task() -> None: +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_falsy_return_from_task(request: pytest.FixtureRequest, checkpointer_name: str): """Test with a falsy return from a task.""" - checkpointer = MemorySaver() + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") @task def falsy_task() -> bool: @@ -5276,17 +5446,18 @@ def test_falsy_return_from_task() -> None: falsy_task().result() interrupt("test") - configurable = {"configurable": {"thread_id": uuid.uuid4()}} + configurable = {"configurable": {"thread_id": str(uuid.uuid4())}} graph.invoke({"a": 5}, configurable) graph.invoke(Command(resume="123"), configurable) -def test_multiple_interrupts_imperative() -> None: +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_multiple_interrupts_imperative( + request: pytest.FixtureRequest, checkpointer_name: str +): """Test multiple interrupts with an imperative API.""" - from langgraph.checkpoint.memory import MemorySaver - from langgraph.func import entrypoint, task + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") - checkpointer = MemorySaver() counter = 0 @task @@ -5307,7 +5478,7 @@ def test_multiple_interrupts_imperative() -> None: return {"values": values} - configurable = {"configurable": {"thread_id": uuid.uuid4()}} + configurable = {"configurable": {"thread_id": str(uuid.uuid4())}} graph.invoke({}, configurable) graph.invoke(Command(resume="a"), configurable) graph.invoke(Command(resume="b"), configurable) @@ -5317,3 +5488,312 @@ def test_multiple_interrupts_imperative() -> None: "values": [2, "a", 4, "b", 6, "c"], } assert counter == 3 + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_double_interrupt_subgraph( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + class AgentState(TypedDict): + input: str + + def node_1(state: AgentState): + result = interrupt("interrupt node 1") + return {"input": result} + + def node_2(state: AgentState): + result = interrupt("interrupt node 2") + return {"input": result} + + subgraph_builder = ( + StateGraph(AgentState) + .add_node("node_1", node_1) + .add_node("node_2", node_2) + .add_edge(START, "node_1") + .add_edge("node_1", "node_2") + .add_edge("node_2", END) + ) + + # invoke the sub graph + subgraph = subgraph_builder.compile(checkpointer=checkpointer) + thread = {"configurable": {"thread_id": str(uuid.uuid4())}} + assert [c for c in subgraph.stream({"input": "test"}, thread)] == [ + { + "__interrupt__": ( + Interrupt( + value="interrupt node 1", + resumable=True, + ns=[AnyStr("node_1:")], + when="during", + ), + ) + }, + ] + # resume from the first interrupt + assert [c for c in subgraph.stream(Command(resume="123"), thread)] == [ + { + "node_1": {"input": "123"}, + }, + { + "__interrupt__": ( + Interrupt( + value="interrupt node 2", + resumable=True, + ns=[AnyStr("node_2:")], + when="during", + ), + ) + }, + ] + # resume from the second interrupt + assert [c for c in subgraph.stream(Command(resume="123"), thread)] == [ + { + "node_2": {"input": "123"}, + }, + ] + + subgraph = subgraph_builder.compile() + + def invoke_sub_agent(state: AgentState): + return subgraph.invoke(state) + + parent_agent = ( + StateGraph(AgentState) + .add_node("invoke_sub_agent", invoke_sub_agent) + .add_edge(START, "invoke_sub_agent") + .add_edge("invoke_sub_agent", END) + .compile(checkpointer=checkpointer) + ) + + assert [c for c in parent_agent.stream({"input": "test"}, thread)] == [ + { + "__interrupt__": ( + Interrupt( + value="interrupt node 1", + resumable=True, + ns=[AnyStr("invoke_sub_agent:"), AnyStr("node_1:")], + when="during", + ), + ) + }, + ] + + # resume from the first interrupt + assert [c for c in parent_agent.stream(Command(resume=True), thread)] == [ + { + "__interrupt__": ( + Interrupt( + value="interrupt node 2", + resumable=True, + ns=[AnyStr("invoke_sub_agent:"), AnyStr("node_2:")], + when="during", + ), + ) + } + ] + + # resume from 2nd interrupt + assert [c for c in parent_agent.stream(Command(resume=True), thread)] == [ + { + "invoke_sub_agent": {"input": True}, + }, + ] + + +def test_sync_streaming_with_functional_api() -> None: + """Test streaming with functional API. + + This test verifies that we're able to stream results as they're being generated + rather than have all the results arrive at once after the graph has completed. + + The time of arrival between the two updates corresponding to the two `slow` tasks + should be greater than the time delay between the two tasks. + """ + + time_delay = 0.01 + + @task() + def slow() -> dict: + time.sleep(time_delay) # Simulate a delay of 10 ms + return {"tic": time.time()} + + @entrypoint() + def graph(inputs: dict) -> list: + first = slow().result() + second = slow().result() + return [first, second] + + arrival_times = [] + + for chunk in graph.stream({}): + if "slow" not in chunk: # We'll just look at the updates from `slow` + continue + arrival_times.append(time.time()) + + assert len(arrival_times) == 2 + delta = arrival_times[1] - arrival_times[0] + # Delta cannot be less than 10 ms if it is streaming as results are generated. + assert delta > time_delay + + +def test_entrypoint_without_checkpointer() -> None: + """Test no checkpointer.""" + states = [] + config = {"configurable": {"thread_id": "1"}} + + # Test without previous + @entrypoint() + def foo(inputs: Any) -> Any: + states.append(inputs) + return inputs + + assert foo.invoke({"a": "1"}, config) == {"a": "1"} + + @entrypoint() + def foo(inputs: Any, *, previous: Any) -> Any: + states.append(previous) + return {"previous": previous, "current": inputs} + + assert foo.invoke({"a": "1"}, config) == {"current": {"a": "1"}, "previous": None} + assert foo.invoke({"a": "1"}, config) == {"current": {"a": "1"}, "previous": None} + + +async def test_async_entrypoint_without_checkpointer() -> None: + """Test no checkpointer.""" + states = [] + config = {"configurable": {"thread_id": "1"}} + + # Test without previous + @entrypoint() + async def foo(inputs: Any) -> Any: + states.append(inputs) + return inputs + + assert (await foo.ainvoke({"a": "1"}, config)) == {"a": "1"} + + @entrypoint() + async def foo(inputs: Any, *, previous: Any) -> Any: + states.append(previous) + return {"previous": previous, "current": inputs} + + assert (await foo.ainvoke({"a": "1"}, config)) == { + "current": {"a": "1"}, + "previous": None, + } + assert (await foo.ainvoke({"a": "1"}, config)) == { + "current": {"a": "1"}, + "previous": None, + } + + +def test_entrypoint_stateful() -> None: + """Test stateful entrypoint invoke.""" + + # Test invoke + states = [] + + @entrypoint(checkpointer=MemorySaver()) + def foo(inputs, *, previous: Any) -> Any: + states.append(previous) + return {"previous": previous, "current": inputs} + + config = {"configurable": {"thread_id": "1"}} + + assert foo.invoke({"a": "1"}, config) == {"current": {"a": "1"}, "previous": None} + assert foo.invoke({"a": "2"}, config) == { + "current": {"a": "2"}, + "previous": {"current": {"a": "1"}, "previous": None}, + } + assert foo.invoke({"a": "3"}, config) == { + "current": {"a": "3"}, + "previous": { + "current": {"a": "2"}, + "previous": {"current": {"a": "1"}, "previous": None}, + }, + } + assert states == [ + None, + {"current": {"a": "1"}, "previous": None}, + {"current": {"a": "2"}, "previous": {"current": {"a": "1"}, "previous": None}}, + ] + + # Test stream + @entrypoint(checkpointer=MemorySaver()) + def foo(inputs, *, previous: Any) -> Any: + return {"previous": previous, "current": inputs} + + config = {"configurable": {"thread_id": "1"}} + items = [item for item in foo.stream({"a": "1"}, config)] + assert items == [{"foo": {"current": {"a": "1"}, "previous": None}}] + + +def test_entrypoint_from_sync_generator() -> None: + """@entrypoint does not support sync generators.""" + previous_return_values = [] + + @entrypoint(checkpointer=MemorySaver()) + def foo(inputs, previous=None) -> Any: + previous_return_values.append(previous) + yield "a" + yield "b" + + config = {"configurable": {"thread_id": "1"}} + + assert foo.invoke({"a": "1"}, config) == ["a", "b"] + assert previous_return_values == [None] + assert foo.invoke({"a": "2"}, config) == ["a", "b"] + assert previous_return_values == [None, ["a", "b"]] + + +def test_entrypoint_request_stream_writer() -> None: + """Test using a stream writer with an entrypoint.""" + + @entrypoint(checkpointer=MemorySaver()) + def foo(inputs, writer: StreamWriter) -> Any: + writer("a") + yield "b" + + config = {"configurable": {"thread_id": "1"}} + + # Different invocations + # Are any of these confusing or unexpected? + assert list(foo.invoke({}, config)) == ["b"] + assert list(foo.stream({}, config)) == ["a", "b"] + + # Stream modes + assert list(foo.stream({}, config, stream_mode=["updates"])) == [ + ("updates", {"foo": ["b"]}) + ] + assert list(foo.stream({}, config, stream_mode=["values"])) == [("values", ["b"])] + assert list(foo.stream({}, config, stream_mode=["custom"])) == [ + ( + "custom", + "a", + ), + ( + "custom", + "b", + ), + ] + + +async def test_entrypoint_from_async_generator() -> None: + """@entrypoint does not support sync generators.""" + # Test invoke + previous_return_values = [] + + # In this version reducers do not work + @entrypoint(checkpointer=MemorySaver()) + async def foo(inputs, previous=None) -> Any: + previous_return_values.append(previous) + yield "a" + yield "b" + + config = {"configurable": {"thread_id": "1"}} + + assert list(await foo.ainvoke({"a": "1"}, config)) == ["a", "b"] + assert previous_return_values == [None] + assert list(foo.invoke({"a": "2"}, config)) == ["a", "b"] + assert previous_return_values == [None, ["a", "b"]] diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index c05682f7e..6767a1510 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -48,7 +48,7 @@ from langgraph.checkpoint.base import ( ) from langgraph.checkpoint.memory import MemorySaver from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, PUSH, START -from langgraph.errors import InvalidUpdateError, MultipleSubgraphsError, NodeInterrupt +from langgraph.errors import InvalidUpdateError, NodeInterrupt from langgraph.func import entrypoint, task from langgraph.graph import END, Graph, StateGraph from langgraph.graph.message import MessagesState, add_messages @@ -89,6 +89,11 @@ logger = logging.getLogger(__name__) pytestmark = pytest.mark.anyio +NEEDS_CONTEXTVARS = pytest.mark.skipif( + sys.version_info < (3, 11), + reason="Python 3.11+ is required for async contextvars support", +) + async def test_checkpoint_errors() -> None: class FaultyGetCheckpointer(MemorySaver): @@ -501,10 +506,7 @@ async def test_node_cancellation_on_other_node_exception_two() -> None: await graph.ainvoke(1) -@pytest.mark.skipif( - sys.version_info < (3, 11), - reason="Python 3.11+ is required for async contextvars support", -) +@NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_dynamic_interrupt(checkpointer_name: str) -> None: class State(TypedDict): @@ -678,10 +680,7 @@ async def test_dynamic_interrupt(checkpointer_name: str) -> None: ) -@pytest.mark.skipif( - sys.version_info < (3, 11), - reason="Python 3.11+ is required for async contextvars support", -) +@NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_dynamic_interrupt_subgraph(checkpointer_name: str) -> None: class SubgraphState(TypedDict): @@ -872,10 +871,7 @@ async def test_dynamic_interrupt_subgraph(checkpointer_name: str) -> None: ) -@pytest.mark.skipif( - sys.version_info < (3, 11), - reason="Python 3.11+ is required for async contextvars support", -) +@NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_copy_checkpoint(checkpointer_name: str) -> None: class State(TypedDict): @@ -1079,10 +1075,7 @@ async def test_copy_checkpoint(checkpointer_name: str) -> None: ) -@pytest.mark.skipif( - sys.version_info < (3, 11), - reason="Python 3.11+ is required for async contextvars support", -) +@NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_node_not_cancelled_on_other_node_interrupted( checkpointer_name: str, @@ -2442,10 +2435,7 @@ async def test_send_sequences(checkpointer_name: str) -> None: ] -@pytest.mark.skipif( - sys.version_info < (3, 11), - reason="Python 3.11+ is required for async contextvars support", -) +@NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_imp_task(checkpointer_name: str) -> None: async with awith_checkpointer(checkpointer_name) as checkpointer: @@ -2493,10 +2483,7 @@ async def test_imp_task(checkpointer_name: str) -> None: assert mapper_calls == 2 -@pytest.mark.skipif( - sys.version_info < (3, 11), - reason="Python 3.11+ is required for async contextvars support", -) +@NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_imp_task_cancel(checkpointer_name: str) -> None: async with awith_checkpointer(checkpointer_name) as checkpointer: @@ -2547,10 +2534,7 @@ async def test_imp_task_cancel(checkpointer_name: str) -> None: assert mapper_cancels == 2 -@pytest.mark.skipif( - sys.version_info < (3, 11), - reason="Python 3.11+ is required for async contextvars support", -) +@NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_imp_sync_from_async(checkpointer_name: str) -> None: async with awith_checkpointer(checkpointer_name) as checkpointer: @@ -2583,10 +2567,7 @@ async def test_imp_sync_from_async(checkpointer_name: str) -> None: ] -@pytest.mark.skipif( - sys.version_info < (3, 11), - reason="Python 3.11+ is required for async contextvars support", -) +@NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_imp_stream_order(checkpointer_name: str) -> None: async with awith_checkpointer(checkpointer_name) as checkpointer: @@ -4068,9 +4049,8 @@ async def test_invoke_join_then_call_other_pregel( async with awith_checkpointer(checkpointer_name) as checkpointer: # add checkpointer app.checkpointer = checkpointer - # subgraph is called twice in the same node, through .map(), so raises - with pytest.raises(MultipleSubgraphsError): - await app.ainvoke([2, 3], {"configurable": {"thread_id": "1"}}) + # subgraph is called twice, and that works + assert await app.ainvoke([2, 3], {"configurable": {"thread_id": "1"}}) == 27 # set inner graph checkpointer NeverCheckpoint inner_app.checkpointer = False @@ -4294,10 +4274,10 @@ async def test_in_one_fan_out_state_graph_waiting_edge(checkpointer_name: str) - docs: Annotated[list[str], sorted_add] async def rewrite_query(data: State) -> State: - return {"query": f'query: {data["query"]}'} + return {"query": f"query: {data['query']}"} async def analyzer_one(data: State) -> State: - return {"query": f'analyzed: {data["query"]}'} + return {"query": f"analyzed: {data['query']}"} async def retriever_one(data: State) -> State: return {"docs": ["doc1", "doc2"]} @@ -4384,10 +4364,10 @@ async def test_in_one_fan_out_state_graph_waiting_edge_via_branch( docs: Annotated[list[str], sorted_add] async def rewrite_query(data: State) -> State: - return {"query": f'query: {data["query"]}'} + return {"query": f"query: {data['query']}"} async def analyzer_one(data: State) -> State: - return {"query": f'analyzed: {data["query"]}'} + return {"query": f"analyzed: {data['query']}"} async def retriever_one(data: State) -> State: return {"docs": ["doc1", "doc2"]} @@ -4801,11 +4781,11 @@ async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular( docs: Annotated[list[str], sorted_add] async def rewrite_query(data: State) -> State: - return {"query": f'query: {data["query"]}'} + return {"query": f"query: {data['query']}"} async def analyzer_one(data: State) -> State: await asyncio.sleep(0.1) - return {"query": f'analyzed: {data["query"]}'} + return {"query": f"analyzed: {data['query']}"} async def retriever_one(data: State) -> State: return {"docs": ["doc1", "doc2"]} @@ -4895,10 +4875,10 @@ async def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None: docs: Annotated[list[str], sorted_add] async def rewrite_query(data: State) -> State: - return {"query": f'query: {data["query"]}'} + return {"query": f"query: {data['query']}"} async def analyzer_one(data: State) -> State: - return {"query": f'analyzed: {data["query"]}'} + return {"query": f"analyzed: {data['query']}"} async def retriever_one(data: State) -> State: return {"docs": ["doc1", "doc2"]} @@ -4979,13 +4959,13 @@ async def test_in_one_fan_out_state_graph_waiting_edge_multiple_cond_edge() -> N docs: Annotated[list[str], sorted_add] async def rewrite_query(data: State) -> State: - return {"query": f'query: {data["query"]}'} + return {"query": f"query: {data['query']}"} async def retriever_picker(data: State) -> list[str]: return ["analyzer_one", "retriever_two"] async def analyzer_one(data: State) -> State: - return {"query": f'analyzed: {data["query"]}'} + return {"query": f"analyzed: {data['query']}"} async def retriever_one(data: State) -> State: return {"docs": ["doc1", "doc2"]} @@ -6117,10 +6097,7 @@ async def test_parent_command(checkpointer_name: str) -> None: ) -@pytest.mark.skipif( - sys.version_info < (3, 11), - reason="Python 3.11+ is required for async contextvars support", -) +@NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_interrupt_subgraph(checkpointer_name: str): class State(TypedDict): @@ -6153,10 +6130,7 @@ async def test_interrupt_subgraph(checkpointer_name: str): assert await graph.ainvoke(Command(resume="bar"), thread1) -@pytest.mark.skipif( - sys.version_info < (3, 11), - reason="Python 3.11+ is required for async contextvars support", -) +@NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_interrupt_multiple(checkpointer_name: str): class State(TypedDict): @@ -6220,10 +6194,7 @@ async def test_interrupt_multiple(checkpointer_name: str): ] -@pytest.mark.skipif( - sys.version_info < (3, 11), - reason="Python 3.11+ is required for async contextvars support", -) +@NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_interrupt_loop(checkpointer_name: str): class State(TypedDict): @@ -6310,6 +6281,63 @@ async def test_interrupt_loop(checkpointer_name: str): ] +@NEEDS_CONTEXTVARS +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_interrupt_functional(checkpointer_name: str) -> None: + @task + async def foo(state: dict) -> dict: + return {"a": state["a"] + "foo"} + + @task + async def bar(state: dict) -> dict: + return {"a": state["a"] + "bar", "b": state["b"]} + + async with awith_checkpointer(checkpointer_name) as checkpointer: + + @entrypoint(checkpointer=checkpointer) + async def graph(inputs: dict) -> dict: + foo_result = await foo(inputs) + value = interrupt("Provide value for bar:") + bar_input = {**foo_result, "b": value} + bar_result = await bar(bar_input) + return bar_result + + config = {"configurable": {"thread_id": "1"}} + # First run, interrupted at bar + await graph.ainvoke({"a": ""}, config) + # Resume with an answer + res = await graph.ainvoke(Command(resume="bar"), config) + assert res == {"a": "foobar", "b": "bar"} + + +@NEEDS_CONTEXTVARS +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_interrupt_task_functional(checkpointer_name: str) -> None: + @task + async def foo(state: dict) -> dict: + return {"a": state["a"] + "foo"} + + @task + async def bar(state: dict) -> dict: + value = interrupt("Provide value for bar:") + return {"a": state["a"] + value} + + async with awith_checkpointer(checkpointer_name) as checkpointer: + + @entrypoint(checkpointer=checkpointer) + async def graph(inputs: dict) -> dict: + foo_result = await foo(inputs) + bar_result = await bar(foo_result) + return bar_result + + config = {"configurable": {"thread_id": "1"}} + # First run, interrupted at bar + await graph.ainvoke({"a": ""}, config) + # Resume with an answer + res = await graph.ainvoke(Command(resume="bar"), config) + assert res == {"a": "foobar"} + + @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_command_with_static_breakpoints(checkpointer_name: str) -> None: """Test that we can use Command to resume and update with static breakpoints.""" @@ -6508,10 +6536,7 @@ async def test_parallel_node_execution(): assert duration < 3.0 -@pytest.mark.skipif( - sys.version_info < (3, 11), - reason="Python 3.11+ is required for async contextvars support", -) +@NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_multiple_interrupt_state_persistence(checkpointer_name: str) -> None: """Test that state is preserved correctly across multiple interrupts.""" @@ -6692,33 +6717,29 @@ async def test_multiple_updates() -> None: ] -@pytest.mark.skipif( - sys.version_info < (3, 11), - reason="Python 3.11+ is required for async contextvars support", -) -async def test_falsy_return_from_task() -> None: +@NEEDS_CONTEXTVARS +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_falsy_return_from_task(checkpointer_name: str) -> None: """Test with a falsy return from a task.""" - checkpointer = MemorySaver() @task async def falsy_task() -> bool: return False - @entrypoint(checkpointer=checkpointer) - async def graph(state: dict) -> dict: - """React tool.""" - await falsy_task() - interrupt("test") + async with awith_checkpointer(checkpointer_name) as checkpointer: - configurable = {"configurable": {"thread_id": uuid.uuid4()}} - await graph.ainvoke({"a": 5}, configurable) - await graph.ainvoke(Command(resume="123"), configurable) + @entrypoint(checkpointer=checkpointer) + async def graph(state: dict) -> dict: + """React tool.""" + await falsy_task() + interrupt("test") + + configurable = {"configurable": {"thread_id": str(uuid.uuid4())}} + await graph.ainvoke({"a": 5}, configurable) + await graph.ainvoke(Command(resume="123"), configurable) -@pytest.mark.skipif( - sys.version_info < (3, 11), - reason="Python 3.11+ is required for async contextvars support", -) +@NEEDS_CONTEXTVARS @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_multiple_interrupts_imperative(checkpointer_name: str) -> None: """Test multiple interrupts with an imperative API.""" @@ -6756,3 +6777,154 @@ async def test_multiple_interrupts_imperative(checkpointer_name: str) -> None: "values": [2, "a", 4, "b", 6, "c"], } assert counter == 3 + + +@NEEDS_CONTEXTVARS +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_double_interrupt_subgraph(checkpointer_name: str) -> None: + class AgentState(TypedDict): + input: str + + def node_1(state: AgentState): + result = interrupt("interrupt node 1") + return {"input": result} + + def node_2(state: AgentState): + result = interrupt("interrupt node 2") + return {"input": result} + + subgraph_builder = ( + StateGraph(AgentState) + .add_node("node_1", node_1) + .add_node("node_2", node_2) + .add_edge(START, "node_1") + .add_edge("node_1", "node_2") + .add_edge("node_2", END) + ) + + async with awith_checkpointer(checkpointer_name) as checkpointer: + # invoke the sub graph + subgraph = subgraph_builder.compile(checkpointer=checkpointer) + thread = {"configurable": {"thread_id": str(uuid.uuid4())}} + assert [c async for c in subgraph.astream({"input": "test"}, thread)] == [ + { + "__interrupt__": ( + Interrupt( + value="interrupt node 1", + resumable=True, + ns=[AnyStr("node_1:")], + when="during", + ), + ) + }, + ] + # resume from the first interrupt + assert [c async for c in subgraph.astream(Command(resume="123"), thread)] == [ + { + "node_1": {"input": "123"}, + }, + { + "__interrupt__": ( + Interrupt( + value="interrupt node 2", + resumable=True, + ns=[AnyStr("node_2:")], + when="during", + ), + ) + }, + ] + # resume from the second interrupt + assert [c async for c in subgraph.astream(Command(resume="123"), thread)] == [ + { + "node_2": {"input": "123"}, + }, + ] + + subgraph = subgraph_builder.compile() + + def invoke_sub_agent(state: AgentState): + return subgraph.invoke(state) + + parent_agent = ( + StateGraph(AgentState) + .add_node("invoke_sub_agent", invoke_sub_agent) + .add_edge(START, "invoke_sub_agent") + .add_edge("invoke_sub_agent", END) + .compile(checkpointer=checkpointer) + ) + + assert [c async for c in parent_agent.astream({"input": "test"}, thread)] == [ + { + "__interrupt__": ( + Interrupt( + value="interrupt node 1", + resumable=True, + ns=[AnyStr("invoke_sub_agent:"), AnyStr("node_1:")], + when="during", + ), + ) + }, + ] + + # resume from the first interrupt + assert [ + c async for c in parent_agent.astream(Command(resume=True), thread) + ] == [ + { + "__interrupt__": ( + Interrupt( + value="interrupt node 2", + resumable=True, + ns=[AnyStr("invoke_sub_agent:"), AnyStr("node_2:")], + when="during", + ), + ) + } + ] + + # resume from 2nd interrupt + assert [ + c async for c in parent_agent.astream(Command(resume=True), thread) + ] == [ + { + "invoke_sub_agent": {"input": True}, + }, + ] + + +@NEEDS_CONTEXTVARS +async def test_async_streaming_with_functional_api() -> None: + """Test streaming with functional API. + + This test verifies that we're able to stream results as they're being generated + rather than have all the results arrive at once after the graph has completed. + + The time of arrival between the two updates corresponding to the two `slow` tasks + should be greater than the time delay between the two tasks. + """ + + time_delay = 0.01 + + @task() + async def slow() -> dict: + await asyncio.sleep(time_delay) # Simulate a delay of 10 ms + return {"tic": asyncio.get_running_loop().time()} + + @entrypoint() + async def graph(inputs: dict) -> list: + first = await slow() + second = await slow() + return [first, second] + + arrival_times = [] + + async for chunk in graph.astream({}): + if "slow" not in chunk: # We'll just look at the updates from `slow` + continue + arrival_times.append(asyncio.get_running_loop().time()) + + assert len(arrival_times) == 2 + delta = arrival_times[1] - arrival_times[0] + # Delta cannot be less than 10 ms if it is streaming as results are generated. + assert delta > time_delay diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py index 970d55be8..fa9a221d0 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/executor.py @@ -1,5 +1,6 @@ import asyncio import concurrent.futures +from collections.abc import Sequence from contextlib import ( AbstractAsyncContextManager, AbstractContextManager, @@ -7,7 +8,7 @@ from contextlib import ( ExitStack, ) from functools import partial -from typing import Any, Optional, Sequence +from typing import Any, Optional from uuid import UUID import orjson @@ -15,7 +16,7 @@ from langchain_core.runnables import RunnableConfig from typing_extensions import Self import langgraph.scheduler.kafka.serde as serde -from langgraph.constants import CONFIG_KEY_DELEGATE, ERROR, NS_END, NS_SEP +from langgraph.constants import CONFIG_KEY_DELEGATE, ERROR from langgraph.errors import CheckpointNotLatest, GraphDelegate, TaskNotFound from langgraph.pregel import Pregel from langgraph.pregel.algo import prepare_single_task @@ -39,7 +40,7 @@ from langgraph.scheduler.kafka.types import ( Topics, ) from langgraph.types import LoopProtocol, PregelExecutableTask, RetryPolicy -from langgraph.utils.config import patch_configurable +from langgraph.utils.config import patch_configurable, recast_checkpoint_ns class AsyncKafkaExecutor(AbstractAsyncContextManager): @@ -165,14 +166,12 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager): # find graph if checkpoint_ns := msg["config"]["configurable"].get("checkpoint_ns"): # remove task_ids from checkpoint_ns - recast_checkpoint_ns = NS_SEP.join( - part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP) - ) + recast = recast_checkpoint_ns(checkpoint_ns) # find the subgraph with the matching name - if recast_checkpoint_ns in self.subgraphs: - graph = self.subgraphs[recast_checkpoint_ns] + if recast in self.subgraphs: + graph = self.subgraphs[recast] else: - raise ValueError(f"Subgraph {recast_checkpoint_ns} not found") + raise ValueError(f"Subgraph {recast} not found") else: graph = self.graph # process message @@ -183,16 +182,19 @@ class AsyncKafkaExecutor(AbstractAsyncContextManager): raise RuntimeError("Checkpoint not found") if saved.checkpoint["id"] != msg["config"]["configurable"]["checkpoint_id"]: raise CheckpointNotLatest() - async with AsyncChannelsManager( - graph.channels, - saved.checkpoint, - LoopProtocol( - config=msg["config"], - store=self.graph.store, - step=saved.metadata["step"] + 1, - stop=saved.metadata["step"] + 2, - ), - ) as (channels, managed), AsyncBackgroundExecutor(msg["config"]) as submit: + async with ( + AsyncChannelsManager( + graph.channels, + saved.checkpoint, + LoopProtocol( + config=msg["config"], + store=self.graph.store, + step=saved.metadata["step"] + 1, + stop=saved.metadata["step"] + 2, + ), + ) as (channels, managed), + AsyncBackgroundExecutor(msg["config"]) as submit, + ): if task := await asyncio.to_thread( prepare_single_task, msg["task"]["path"], @@ -378,14 +380,12 @@ class KafkaExecutor(AbstractContextManager): # find graph if checkpoint_ns := msg["config"]["configurable"].get("checkpoint_ns"): # remove task_ids from checkpoint_ns - recast_checkpoint_ns = NS_SEP.join( - part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP) - ) + recast = recast_checkpoint_ns(checkpoint_ns) # find the subgraph with the matching name - if recast_checkpoint_ns in self.subgraphs: - graph = self.subgraphs[recast_checkpoint_ns] + if recast in self.subgraphs: + graph = self.subgraphs[recast] else: - raise ValueError(f"Subgraph {recast_checkpoint_ns} not found") + raise ValueError(f"Subgraph {recast} not found") else: graph = self.graph # process message @@ -396,16 +396,19 @@ class KafkaExecutor(AbstractContextManager): raise RuntimeError("Checkpoint not found") if saved.checkpoint["id"] != msg["config"]["configurable"]["checkpoint_id"]: raise CheckpointNotLatest() - with ChannelsManager( - graph.channels, - saved.checkpoint, - LoopProtocol( - config=msg["config"], - store=self.graph.store, - step=saved.metadata["step"] + 1, - stop=saved.metadata["step"] + 2, - ), - ) as (channels, managed), BackgroundExecutor({}) as submit: + with ( + ChannelsManager( + graph.channels, + saved.checkpoint, + LoopProtocol( + config=msg["config"], + store=self.graph.store, + step=saved.metadata["step"] + 1, + stop=saved.metadata["step"] + 2, + ), + ) as (channels, managed), + BackgroundExecutor({}) as submit, + ): if task := prepare_single_task( msg["task"]["path"], msg["task"]["id"], diff --git a/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py b/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py index 4e5be8470..9ed72cd02 100644 --- a/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py +++ b/libs/scheduler-kafka/langgraph/scheduler/kafka/orchestrator.py @@ -13,11 +13,11 @@ from typing_extensions import Self import langgraph.scheduler.kafka.serde as serde from langgraph.constants import ( + CONF, CONFIG_KEY_DEDUPE_TASKS, CONFIG_KEY_ENSURE_LATEST, + CONFIG_KEY_SCRATCHPAD, INTERRUPT, - NS_END, - NS_SEP, SCHEDULED, ) from langgraph.errors import CheckpointNotLatest, GraphInterrupt @@ -37,7 +37,7 @@ from langgraph.scheduler.kafka.types import ( Topics, ) from langgraph.types import RetryPolicy -from langgraph.utils.config import patch_configurable +from langgraph.utils.config import patch_configurable, recast_checkpoint_ns class AsyncKafkaOrchestrator(AbstractAsyncContextManager): @@ -140,14 +140,12 @@ class AsyncKafkaOrchestrator(AbstractAsyncContextManager): # find graph if checkpoint_ns := msg["config"]["configurable"].get("checkpoint_ns"): # remove task_ids from checkpoint_ns - recast_checkpoint_ns = NS_SEP.join( - part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP) - ) + recast = recast_checkpoint_ns(checkpoint_ns) # find the subgraph with the matching name - if recast_checkpoint_ns in self.subgraphs: - graph = self.subgraphs[recast_checkpoint_ns] + if recast in self.subgraphs: + graph = self.subgraphs[recast] else: - raise ValueError(f"Subgraph {recast_checkpoint_ns} not found") + raise ValueError(f"Subgraph {recast} not found") else: graph = self.graph # process message @@ -163,7 +161,6 @@ class AsyncKafkaOrchestrator(AbstractAsyncContextManager): stream_keys=graph.stream_channels, interrupt_after=graph.interrupt_after_nodes, interrupt_before=graph.interrupt_before_nodes, - check_subgraphs=False, ) as loop: if loop.tick(input_keys=graph.input_channels): # wait for checkpoint to be saved @@ -173,6 +170,16 @@ class AsyncKafkaOrchestrator(AbstractAsyncContextManager): if new_tasks := [ t for t in loop.tasks.values() if not t.scheduled and not t.writes ]: + config = patch_configurable( + loop.config, + { + **loop.checkpoint_config["configurable"], + CONFIG_KEY_DEDUPE_TASKS: True, + CONFIG_KEY_ENSURE_LATEST: True, + }, + ) + if CONFIG_KEY_SCRATCHPAD in config[CONF]: + config[CONF][CONFIG_KEY_SCRATCHPAD]["subgraph_counter"] = 0 # send messages to executor futures = await asyncio.gather( *( @@ -180,16 +187,7 @@ class AsyncKafkaOrchestrator(AbstractAsyncContextManager): self.topics.executor, value=serde.dumps( MessageToExecutor( - config=patch_configurable( - loop.config, - { - **loop.checkpoint_config[ - "configurable" - ], - CONFIG_KEY_DEDUPE_TASKS: True, - CONFIG_KEY_ENSURE_LATEST: True, - }, - ), + config=config, task=ExecutorTask(id=task.id, path=task.path), finally_send=msg.get("finally_send"), ) @@ -330,14 +328,12 @@ class KafkaOrchestrator(AbstractContextManager): # find graph if checkpoint_ns := msg["config"]["configurable"].get("checkpoint_ns"): # remove task_ids from checkpoint_ns - recast_checkpoint_ns = NS_SEP.join( - part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP) - ) + recast = recast_checkpoint_ns(checkpoint_ns) # find the subgraph with the matching name - if recast_checkpoint_ns in self.subgraphs: - graph = self.subgraphs[recast_checkpoint_ns] + if recast in self.subgraphs: + graph = self.subgraphs[recast] else: - raise ValueError(f"Subgraph {recast_checkpoint_ns} not found") + raise ValueError(f"Subgraph {recast} not found") else: graph = self.graph # process message @@ -353,7 +349,6 @@ class KafkaOrchestrator(AbstractContextManager): stream_keys=graph.stream_channels, interrupt_after=graph.interrupt_after_nodes, interrupt_before=graph.interrupt_before_nodes, - check_subgraphs=False, ) as loop: if loop.tick(input_keys=graph.input_channels): # wait for checkpoint to be saved @@ -363,20 +358,23 @@ class KafkaOrchestrator(AbstractContextManager): if new_tasks := [ t for t in loop.tasks.values() if not t.scheduled and not t.writes ]: + config = patch_configurable( + loop.config, + { + **loop.checkpoint_config["configurable"], + CONFIG_KEY_DEDUPE_TASKS: True, + CONFIG_KEY_ENSURE_LATEST: True, + }, + ) + if CONFIG_KEY_SCRATCHPAD in config[CONF]: + config[CONF][CONFIG_KEY_SCRATCHPAD]["subgraph_counter"] = 0 # send messages to executor futures = [ self.producer.send( self.topics.executor, value=serde.dumps( MessageToExecutor( - config=patch_configurable( - loop.config, - { - **loop.checkpoint_config["configurable"], - CONFIG_KEY_DEDUPE_TASKS: True, - CONFIG_KEY_ENSURE_LATEST: True, - }, - ), + config=config, task=ExecutorTask(id=task.id, path=task.path), finally_send=msg.get("finally_send"), ) diff --git a/libs/scheduler-kafka/tests/any.py b/libs/scheduler-kafka/tests/any.py index 3ea224173..0336d8506 100644 --- a/libs/scheduler-kafka/tests/any.py +++ b/libs/scheduler-kafka/tests/any.py @@ -53,3 +53,11 @@ class AnyList(list): return False else: return True + + +class AnyInt(int): + def __init__(self) -> None: + super().__init__() + + def __eq__(self, other: object) -> bool: + return isinstance(other, int) diff --git a/libs/scheduler-kafka/tests/test_subgraph.py b/libs/scheduler-kafka/tests/test_subgraph.py index 1e1f1e396..2a6c9992a 100644 --- a/libs/scheduler-kafka/tests/test_subgraph.py +++ b/libs/scheduler-kafka/tests/test_subgraph.py @@ -15,7 +15,7 @@ from langgraph.graph.state import StateGraph from langgraph.pregel import Pregel from langgraph.scheduler.kafka import serde from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics -from tests.any import AnyDict, AnyList +from tests.any import AnyDict, AnyInt from tests.drain import drain_topics_async from tests.messages import _AnyIdAIMessage, _AnyIdHumanMessage @@ -195,10 +195,16 @@ async def test_subgraph_w_interrupt( "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, "__pregel_resuming": False, + "__pregel_previous": None, "__pregel_store": None, "__pregel_task_id": history[0].tasks[0].id, - "__pregel_scratchpad": {}, - "__pregel_writes": AnyList(), + "__pregel_scratchpad": { + "subgraph_counter": AnyInt(), + "call_counter": 0, + "interrupt_counter": -1, + "null_resume": None, + "resume": [], + }, "checkpoint_id": None, "checkpoint_map": { "": history[0].config["configurable"]["checkpoint_id"] @@ -262,10 +268,16 @@ async def test_subgraph_w_interrupt( "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, "__pregel_resuming": False, + "__pregel_previous": None, "__pregel_store": None, "__pregel_task_id": history[0].tasks[0].id, - "__pregel_scratchpad": {}, - "__pregel_writes": AnyList(), + "__pregel_scratchpad": { + "subgraph_counter": AnyInt(), + "call_counter": 0, + "interrupt_counter": -1, + "null_resume": None, + "resume": [], + }, "checkpoint_id": c.config["configurable"]["checkpoint_id"], "checkpoint_map": { "": history[0].config["configurable"]["checkpoint_id"] @@ -359,10 +371,16 @@ async def test_subgraph_w_interrupt( "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, "__pregel_resuming": False, + "__pregel_previous": None, "__pregel_store": None, "__pregel_task_id": history[0].tasks[0].id, - "__pregel_scratchpad": {}, - "__pregel_writes": AnyList(), + "__pregel_scratchpad": { + "subgraph_counter": AnyInt(), + "call_counter": 0, + "interrupt_counter": -1, + "null_resume": None, + "resume": [], + }, "checkpoint_id": c.config["configurable"]["checkpoint_id"], "checkpoint_map": { "": history[0].config["configurable"]["checkpoint_id"] @@ -466,10 +484,16 @@ async def test_subgraph_w_interrupt( "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, "__pregel_resuming": True, + "__pregel_previous": None, "__pregel_store": None, "__pregel_task_id": history[1].tasks[0].id, - "__pregel_scratchpad": {}, - "__pregel_writes": AnyList(), + "__pregel_scratchpad": { + "subgraph_counter": AnyInt(), + "call_counter": 0, + "interrupt_counter": -1, + "null_resume": None, + "resume": [], + }, "checkpoint_id": None, "checkpoint_map": { "": history[1].config["configurable"]["checkpoint_id"] @@ -528,10 +552,16 @@ async def test_subgraph_w_interrupt( "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, "__pregel_resuming": True, + "__pregel_previous": None, "__pregel_store": None, "__pregel_task_id": history[1].tasks[0].id, - "__pregel_scratchpad": {}, - "__pregel_writes": AnyList(), + "__pregel_scratchpad": { + "subgraph_counter": AnyInt(), + "call_counter": 0, + "interrupt_counter": -1, + "null_resume": None, + "resume": [], + }, "checkpoint_id": c.config["configurable"]["checkpoint_id"], "checkpoint_map": { "": history[1].config["configurable"]["checkpoint_id"] @@ -646,10 +676,16 @@ async def test_subgraph_w_interrupt( "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, "__pregel_resuming": True, + "__pregel_previous": None, "__pregel_store": None, "__pregel_task_id": history[1].tasks[0].id, - "__pregel_scratchpad": {}, - "__pregel_writes": AnyList(), + "__pregel_scratchpad": { + "subgraph_counter": AnyInt(), + "call_counter": 0, + "interrupt_counter": -1, + "null_resume": None, + "resume": [], + }, "checkpoint_id": c.config["configurable"]["checkpoint_id"], "checkpoint_map": { "": history[1].config["configurable"]["checkpoint_id"] diff --git a/libs/scheduler-kafka/tests/test_subgraph_sync.py b/libs/scheduler-kafka/tests/test_subgraph_sync.py index 210312b3b..c2c9a8fc1 100644 --- a/libs/scheduler-kafka/tests/test_subgraph_sync.py +++ b/libs/scheduler-kafka/tests/test_subgraph_sync.py @@ -15,7 +15,7 @@ from langgraph.pregel import Pregel from langgraph.scheduler.kafka import serde from langgraph.scheduler.kafka.default_sync import DefaultProducer from langgraph.scheduler.kafka.types import MessageToOrchestrator, Topics -from tests.any import AnyDict, AnyList +from tests.any import AnyDict, AnyInt from tests.drain import drain_topics from tests.messages import _AnyIdAIMessage, _AnyIdHumanMessage @@ -194,10 +194,16 @@ def test_subgraph_w_interrupt( "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, "__pregel_resuming": False, + "__pregel_previous": None, "__pregel_store": None, "__pregel_task_id": history[0].tasks[0].id, - "__pregel_scratchpad": {}, - "__pregel_writes": AnyList(), + "__pregel_scratchpad": { + "subgraph_counter": AnyInt(), + "call_counter": 0, + "interrupt_counter": -1, + "null_resume": None, + "resume": [], + }, "checkpoint_id": None, "checkpoint_map": { "": history[0].config["configurable"]["checkpoint_id"] @@ -262,9 +268,15 @@ def test_subgraph_w_interrupt( "__pregel_store": None, "__pregel_dedupe_tasks": True, "__pregel_resuming": False, + "__pregel_previous": None, "__pregel_task_id": history[0].tasks[0].id, - "__pregel_scratchpad": {}, - "__pregel_writes": AnyList(), + "__pregel_scratchpad": { + "subgraph_counter": AnyInt(), + "call_counter": 0, + "interrupt_counter": -1, + "null_resume": None, + "resume": [], + }, "checkpoint_id": c.config["configurable"]["checkpoint_id"], "checkpoint_map": { "": history[0].config["configurable"]["checkpoint_id"] @@ -360,8 +372,14 @@ def test_subgraph_w_interrupt( "__pregel_store": None, "__pregel_resuming": False, "__pregel_task_id": history[0].tasks[0].id, - "__pregel_scratchpad": {}, - "__pregel_writes": AnyList(), + "__pregel_previous": None, + "__pregel_scratchpad": { + "subgraph_counter": AnyInt(), + "call_counter": 0, + "interrupt_counter": -1, + "null_resume": None, + "resume": [], + }, "checkpoint_id": c.config["configurable"]["checkpoint_id"], "checkpoint_map": { "": history[0].config["configurable"]["checkpoint_id"] @@ -465,9 +483,15 @@ def test_subgraph_w_interrupt( "__pregel_dedupe_tasks": True, "__pregel_store": None, "__pregel_resuming": True, + "__pregel_previous": None, "__pregel_task_id": history[1].tasks[0].id, - "__pregel_scratchpad": {}, - "__pregel_writes": AnyList(), + "__pregel_scratchpad": { + "subgraph_counter": AnyInt(), + "call_counter": 0, + "interrupt_counter": -1, + "null_resume": None, + "resume": [], + }, "checkpoint_id": None, "checkpoint_map": { "": history[1].config["configurable"]["checkpoint_id"] @@ -527,9 +551,15 @@ def test_subgraph_w_interrupt( "__pregel_dedupe_tasks": True, "__pregel_store": None, "__pregel_resuming": True, + "__pregel_previous": None, "__pregel_task_id": history[1].tasks[0].id, - "__pregel_scratchpad": {}, - "__pregel_writes": AnyList(), + "__pregel_scratchpad": { + "subgraph_counter": AnyInt(), + "call_counter": 0, + "interrupt_counter": -1, + "null_resume": None, + "resume": [], + }, "checkpoint_id": c.config["configurable"]["checkpoint_id"], "checkpoint_map": { "": history[1].config["configurable"]["checkpoint_id"] @@ -644,10 +674,16 @@ def test_subgraph_w_interrupt( "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, "__pregel_resuming": True, + "__pregel_previous": None, "__pregel_store": None, "__pregel_task_id": history[1].tasks[0].id, - "__pregel_scratchpad": {}, - "__pregel_writes": AnyList(), + "__pregel_scratchpad": { + "subgraph_counter": AnyInt(), + "call_counter": 0, + "interrupt_counter": -1, + "null_resume": None, + "resume": [], + }, "checkpoint_id": c.config["configurable"]["checkpoint_id"], "checkpoint_map": { "": history[1].config["configurable"]["checkpoint_id"]