From 445a1109177c837655c017ffc06c8042f97105f0 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 10 May 2024 13:44:20 -0700 Subject: [PATCH] Add managed values and IsLastStep (#330) * Add managed values and IsLastStep - managed values are read-only state keys whose values are managed by langgraph - this PR implements one: IsLastValue, a boolean which is True in the last iteration, eg to allow you to return a nice "ran out of iterations" message to user * Fix * Fix some issues * py39 * Break ref * Fix types * Fix * Fix * Update state.py * Don't mutate dictionary while using it in for loop. --------- Co-authored-by: Andrew Nguonly --- langgraph/graph/state.py | 46 +++++-- langgraph/managed/__init__.py | 3 + langgraph/managed/base.py | 105 +++++++++++++++ langgraph/managed/is_last_step.py | 12 ++ langgraph/prebuilt/chat_agent_executor.py | 28 +++- langgraph/pregel/__init__.py | 153 +++++++++++++++++++--- langgraph/pregel/read.py | 5 +- tests/__snapshots__/test_pregel.ambr | 4 +- tests/test_pregel.py | 78 ++++++----- 9 files changed, 369 insertions(+), 65 deletions(-) create mode 100644 langgraph/managed/__init__.py create mode 100644 langgraph/managed/base.py create mode 100644 langgraph/managed/is_last_step.py diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py index 6b180a701..9d121159d 100644 --- a/langgraph/graph/state.py +++ b/langgraph/graph/state.py @@ -1,7 +1,7 @@ import logging from functools import partial from inspect import signature -from typing import Any, Optional, Sequence, Type, Union, get_type_hints +from typing import Any, Optional, Sequence, Type, Union, get_origin, get_type_hints from langchain_core.pydantic_v1 import BaseModel from langchain_core.runnables import Runnable, RunnableConfig @@ -17,6 +17,7 @@ from langgraph.checkpoint import BaseCheckpointSaver from langgraph.constants import TAG_HIDDEN from langgraph.errors import InvalidUpdateError from langgraph.graph.graph import END, START, Branch, CompiledGraph, Graph +from langgraph.managed.base import ManagedValue, is_managed_value from langgraph.pregel.read import ChannelRead, PregelNode from langgraph.pregel.types import All from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry @@ -40,7 +41,7 @@ class StateGraph(Graph): super().__init__() self.schema = state_schema self.config_schema = config_schema - self.channels = _get_channels(state_schema) + self.channels, self.managed = _get_channels(state_schema) if any(isinstance(c, BinaryOperatorAggregate) for c in self.channels.values()): self.support_multiple_edges = True self.waiting_edges: set[tuple[tuple[str, ...], str]] = set() @@ -193,6 +194,8 @@ class CompiledStateGraph(CompiledGraph): return super().get_output_schema(config) def attach_node(self, key: str, node: Optional[Runnable]) -> None: + state_keys = list(self.builder.channels) + def _get_state_key(input: dict, config: RunnableConfig, *, key: str) -> Any: if input is None: return SKIP_WRITE @@ -201,7 +204,6 @@ class CompiledStateGraph(CompiledGraph): else: return input.get(key, SKIP_WRITE) - state_keys = list(self.builder.channels) # state updaters state_write_entries = [ ( @@ -231,11 +233,11 @@ class CompiledStateGraph(CompiledGraph): self.channels[key] = EphemeralValue(Any) self.nodes[key] = PregelNode( triggers=[], - # read state keys + # read state keys and managed values channels=( state_keys if state_keys == ["__root__"] - else {chan: chan for chan in state_keys} + else ({chan: chan for chan in state_keys} | self.builder.managed) ), # coerce state dict to schema class (eg. pydantic model) mapper=( @@ -339,19 +341,32 @@ def _coerce_state(schema: Type[Any], input: dict[str, Any]) -> dict[str, Any]: return schema(**input) -def _get_channels(schema: Type[dict]) -> dict[str, BaseChannel]: +def _get_channels( + schema: Type[dict], +) -> tuple[dict[str, BaseChannel], dict[str, Type[ManagedValue]]]: if not hasattr(schema, "__annotations__"): - return {"__root__": _get_channel(schema)} + return {"__root__": _get_channel(schema, allow_managed=False)}, {} - return { + all_keys = { name: _get_channel(typ) for name, typ in get_type_hints(schema, include_extras=True).items() if name != "__slots__" } + return ( + {k: v for k, v in all_keys.items() if not is_managed_value(v)}, + {k: v for k, v in all_keys.items() if is_managed_value(v)}, + ) -def _get_channel(annotation: Any) -> BaseChannel: - if channel := _is_field_binop(annotation): +def _get_channel( + annotation: Any, *, allow_managed: bool = True +) -> Union[BaseChannel, Type[ManagedValue]]: + if manager := _is_field_managed_value(annotation): + if allow_managed: + return manager + else: + raise ValueError(f"This {annotation} not allowed in this position") + elif channel := _is_field_binop(annotation): return channel return LastValue(annotation) @@ -371,3 +386,14 @@ def _is_field_binop(typ: Type[Any]) -> Optional[BinaryOperatorAggregate]: ): return BinaryOperatorAggregate(typ, meta[0]) return None + + +def _is_field_managed_value(typ: Type[Any]) -> Optional[Type[ManagedValue]]: + if hasattr(typ, "__metadata__"): + meta = typ.__metadata__ + if len(meta) == 1: + decoration = get_origin(meta[0]) or meta[0] + if is_managed_value(decoration): + return decoration + + return None diff --git a/langgraph/managed/__init__.py b/langgraph/managed/__init__.py new file mode 100644 index 000000000..2c101e400 --- /dev/null +++ b/langgraph/managed/__init__.py @@ -0,0 +1,3 @@ +from langgraph.managed.is_last_step import IsLastStep + +__all__ = ["IsLastStep"] diff --git a/langgraph/managed/base.py b/langgraph/managed/base.py new file mode 100644 index 000000000..15b47070b --- /dev/null +++ b/langgraph/managed/base.py @@ -0,0 +1,105 @@ +import asyncio +from abc import ABC, abstractmethod +from contextlib import AsyncExitStack, ExitStack, asynccontextmanager, contextmanager +from inspect import isclass +from typing import ( + TYPE_CHECKING, + Any, + AsyncGenerator, + Generator, + Generic, + Sequence, + Type, + TypeVar, +) + +from langchain_core.runnables import RunnableConfig +from typing_extensions import Self, TypeGuard + +from langgraph.pregel.types import PregelTaskDescription + +if TYPE_CHECKING: + from langgraph.pregel import Pregel + +V = TypeVar("V") + + +class ManagedValue(ABC, Generic[V]): + def __init__(self, config: RunnableConfig, graph: "Pregel") -> None: + self.config = config + self.graph = graph + + @classmethod + @contextmanager + def enter( + cls, config: RunnableConfig, graph: "Pregel" + ) -> Generator[Self, None, None]: + try: + value = cls(config, graph) + yield value + finally: + # because managed value and Pregel have reference to each other + # let's make sure to break the reference on exit + try: + del value + except UnboundLocalError: + pass + + @classmethod + @asynccontextmanager + async def aenter( + cls, config: RunnableConfig, graph: "Pregel" + ) -> AsyncGenerator[Self, None]: + try: + value = cls(config, graph) + yield value + finally: + # because managed value and Pregel have reference to each other + # let's make sure to break the reference on exit + try: + del value + except UnboundLocalError: + pass + + @abstractmethod + def __call__(self, step: int, task: PregelTaskDescription) -> V: + ... + + +def is_managed_value(value: Any) -> TypeGuard[Type[ManagedValue]]: + return isclass(value) and issubclass(value, ManagedValue) + + +@contextmanager +def ManagedValuesManager( + values: Sequence[Type[ManagedValue]], + config: RunnableConfig, + graph: "Pregel", +) -> Generator[Sequence[ManagedValue], None, None]: + with ExitStack() as stack: + unique: list[Type[ManagedValue]] = [] + for value in values: + if value not in unique: + unique.append(value) + + yield [stack.enter_context(value.enter(config, graph)) for value in unique] + + +@asynccontextmanager +async def AsyncManagedValuesManager( + values: Sequence[Type[ManagedValue]], + config: RunnableConfig, + graph: "Pregel", +) -> AsyncGenerator[Sequence[ManagedValue], None]: + async with AsyncExitStack() as stack: + unique: list[Type[ManagedValue]] = [] + for value in values: + if value not in unique: + unique.append(value) + + yield await asyncio.gather( + *( + stack.enter_async_context(value.aenter(config, graph)) + for value in unique + ) + ) diff --git a/langgraph/managed/is_last_step.py b/langgraph/managed/is_last_step.py new file mode 100644 index 000000000..5860a9cbe --- /dev/null +++ b/langgraph/managed/is_last_step.py @@ -0,0 +1,12 @@ +from typing import Annotated + +from langgraph.managed.base import ManagedValue +from langgraph.pregel.types import PregelExecutableTask + + +class IsLastStepManager(ManagedValue[bool]): + def __call__(self, step: int, task: PregelExecutableTask) -> bool: + return step == self.config["recursion_limit"] - 1 + + +IsLastStep = Annotated[bool, IsLastStepManager] diff --git a/langgraph/prebuilt/chat_agent_executor.py b/langgraph/prebuilt/chat_agent_executor.py index ab8db7fb6..13ccb1f72 100644 --- a/langgraph/prebuilt/chat_agent_executor.py +++ b/langgraph/prebuilt/chat_agent_executor.py @@ -2,7 +2,12 @@ import json from typing import Annotated, Callable, Optional, Sequence, TypedDict, Union from langchain_core.language_models import LanguageModelLike -from langchain_core.messages import BaseMessage, FunctionMessage, SystemMessage +from langchain_core.messages import ( + AIMessage, + BaseMessage, + FunctionMessage, + SystemMessage, +) from langchain_core.runnables import Runnable, RunnableLambda from langchain_core.tools import BaseTool from langchain_core.utils.function_calling import convert_to_openai_function @@ -12,6 +17,7 @@ from langgraph.checkpoint import BaseCheckpointSaver from langgraph.graph import END, StateGraph from langgraph.graph.graph import CompiledGraph from langgraph.graph.message import add_messages +from langgraph.managed import IsLastStep from langgraph.prebuilt.tool_executor import ToolExecutor, ToolInvocation from langgraph.prebuilt.tool_node import ToolNode @@ -25,6 +31,8 @@ class AgentState(TypedDict): messages: Annotated[Sequence[BaseMessage], add_messages] + is_last_step: IsLastStep + @deprecated("0.0.44", "create_react_agent") def create_function_calling_executor( @@ -238,12 +246,30 @@ def create_react_agent( def call_model(state: AgentState): messages = state["messages"] response = model_runnable.invoke(messages) + if state["is_last_step"] and response.tool_calls: + return { + "messages": [ + AIMessage( + id=response.id, + content="Sorry, need more steps to process this request.", + ) + ] + } # We return a list, because this will get added to the existing list return {"messages": [response]} async def acall_model(state: AgentState): messages = state["messages"] response = await model_runnable.ainvoke(messages) + if state["is_last_step"] and response.tool_calls: + return { + "messages": [ + AIMessage( + id=response.id, + content="Sorry, need more steps to process this request.", + ) + ] + } # We return a list, because this will get added to the existing list return {"messages": [response]} diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index b718ab693..597f617d5 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -4,6 +4,7 @@ import asyncio import concurrent.futures from collections import defaultdict, deque from functools import partial +from inspect import isclass from typing import ( Any, AsyncIterator, @@ -67,6 +68,12 @@ from langgraph.constants import ( TAG_HIDDEN, ) from langgraph.errors import GraphRecursionError, InvalidUpdateError +from langgraph.managed.base import ( + AsyncManagedValuesManager, + ManagedValue, + ManagedValuesManager, + is_managed_value, +) from langgraph.pregel.debug import ( map_debug_checkpoint, map_debug_task_results, @@ -319,6 +326,16 @@ class Pregel( def stream_channels_asis(self) -> Union[str, Sequence[str]]: return self.stream_channels or [k for k in self.channels] + @property + def managed_values_list(self) -> Sequence[Type[ManagedValue]]: + return [ + v + for node in self.nodes.values() + if isinstance(node.channels, dict) + for v in node.channels.values() + if is_managed_value(v) + ] + def get_state(self, config: RunnableConfig) -> StateSnapshot: """Get the current state of the graph.""" if not self.checkpointer: @@ -326,9 +343,20 @@ class Pregel( saved = self.checkpointer.get_tuple(config) checkpoint = saved.checkpoint if saved else empty_checkpoint() - with ChannelsManager(self.channels, checkpoint) as channels: + config = saved.config if saved else config + with ChannelsManager( + self.channels, checkpoint + ) as channels, ManagedValuesManager( + self.managed_values_list, ensure_config(config), self + ) as managed: _, next_tasks = _prepare_next_tasks( - checkpoint, self.nodes, channels, for_execution=False + checkpoint, + self.nodes, + channels, + managed, + config, + -1, + for_execution=False, ) return StateSnapshot( read_channels(channels, self.stream_channels_asis), @@ -345,9 +373,21 @@ class Pregel( saved = await self.checkpointer.aget_tuple(config) checkpoint = saved.checkpoint if saved else empty_checkpoint() - async with AsyncChannelsManager(self.channels, checkpoint) as channels: + + config = saved.config if saved else config + async with AsyncChannelsManager( + self.channels, checkpoint + ) as channels, AsyncManagedValuesManager( + self.managed_values_list, ensure_config(config), self + ) as managed: _, next_tasks = _prepare_next_tasks( - checkpoint, self.nodes, channels, for_execution=False + checkpoint, + self.nodes, + channels, + managed, + config, + -1, + for_execution=False, ) return StateSnapshot( read_channels(channels, self.stream_channels_asis), @@ -371,9 +411,19 @@ class Pregel( for config, checkpoint, metadata, parent_config in self.checkpointer.list( config, before=before, limit=limit ): - with ChannelsManager(self.channels, checkpoint) as channels: + with ChannelsManager( + self.channels, checkpoint + ) as channels, ManagedValuesManager( + self.managed_values_list, ensure_config(config), self + ) as managed: _, next_tasks = _prepare_next_tasks( - checkpoint, self.nodes, channels, for_execution=False + checkpoint, + self.nodes, + channels, + managed, + config, + -1, + for_execution=False, ) yield StateSnapshot( read_channels(channels, self.stream_channels_asis), @@ -400,9 +450,19 @@ class Pregel( metadata, parent_config, ) in self.checkpointer.alist(config, before=before, limit=limit): - async with AsyncChannelsManager(self.channels, checkpoint) as channels: + async with AsyncChannelsManager( + self.channels, checkpoint + ) as channels, AsyncManagedValuesManager( + self.managed_values_list, ensure_config(config), self + ) as managed: _, next_tasks = _prepare_next_tasks( - checkpoint, self.nodes, channels, for_execution=False + checkpoint, + self.nodes, + channels, + managed, + config, + -1, + for_execution=False, ) yield StateSnapshot( read_channels(channels, self.stream_channels_asis), @@ -662,12 +722,22 @@ class Pregel( # create channels from checkpoint with ChannelsManager( self.channels, checkpoint - ) as channels, get_executor_for_config(config) as executor: + ) as channels, get_executor_for_config( + config + ) as executor, ManagedValuesManager( + self.managed_values_list, config, self + ) as managed: # map inputs to channel updates if input_writes := deque(map_input(input_keys, input)): # discard any unfinished tasks from previous checkpoint checkpoint, _ = _prepare_next_tasks( - checkpoint, processes, channels, for_execution=True + checkpoint, + processes, + channels, + managed, + config, + -1, + for_execution=True, ) # apply input writes _apply_writes(checkpoint, channels, input_writes) @@ -708,7 +778,13 @@ class Pregel( stop = start + config["recursion_limit"] + 1 for step in range(start, stop): next_checkpoint, next_tasks = _prepare_next_tasks( - checkpoint, processes, channels, for_execution=True + checkpoint, + processes, + channels, + managed, + config, + step, + for_execution=True, ) # if no more tasks, we're done @@ -741,7 +817,7 @@ class Pregel( proc, input, patch_config( - merge_configs(config, proc_config), + proc_config, run_name=name, callbacks=run_manager.get_child(f"graph:step:{step}"), configurable={ @@ -942,12 +1018,22 @@ class Pregel( checkpoint_config = saved.config if saved else config start = saved.metadata.get("step", -2) + 1 if saved else -1 # create channels from checkpoint - async with AsyncChannelsManager(self.channels, checkpoint) as channels: + async with AsyncChannelsManager( + self.channels, checkpoint + ) as channels, AsyncManagedValuesManager( + self.managed_values_list, config, self + ) as managed: # map inputs to channel updates if input_writes := deque(map_input(input_keys, input)): # discard any unfinished tasks from previous checkpoint checkpoint, _ = _prepare_next_tasks( - checkpoint, processes, channels, for_execution=True + checkpoint, + processes, + channels, + managed, + config, + -1, + for_execution=True, ) # apply input writes _apply_writes(checkpoint, channels, input_writes) @@ -990,7 +1076,13 @@ class Pregel( stop = start + config["recursion_limit"] + 1 for step in range(start, stop): next_checkpoint, next_tasks = _prepare_next_tasks( - checkpoint, processes, channels, for_execution=True + checkpoint, + processes, + channels, + managed, + config, + step, + for_execution=True, ) # if no more tasks, we're done @@ -1023,7 +1115,7 @@ class Pregel( proc, input, patch_config( - merge_configs(config, proc_config), + proc_config, run_name=name, callbacks=run_manager.get_child(f"graph:step:{step}"), configurable={ @@ -1386,6 +1478,9 @@ def _prepare_next_tasks( checkpoint: Checkpoint, processes: Mapping[str, PregelNode], channels: Mapping[str, BaseChannel], + managed: Sequence[ManagedValue], + config: RunnableConfig, + step: int, for_execution: Literal[False], ) -> tuple[Checkpoint, list[PregelTaskDescription]]: ... @@ -1396,6 +1491,9 @@ def _prepare_next_tasks( checkpoint: Checkpoint, processes: Mapping[str, PregelNode], channels: Mapping[str, BaseChannel], + managed: Sequence[ManagedValue], + config: RunnableConfig, + step: int, for_execution: Literal[True], ) -> tuple[Checkpoint, list[PregelExecutableTask]]: ... @@ -1405,6 +1503,9 @@ def _prepare_next_tasks( checkpoint: Checkpoint, processes: Mapping[str, PregelNode], channels: Mapping[str, BaseChannel], + managed: Sequence[ManagedValue], + config: RunnableConfig, + step: int, *, for_execution: bool, ) -> tuple[Checkpoint, Union[list[PregelTaskDescription], list[PregelExecutableTask]]]: @@ -1427,10 +1528,21 @@ def _prepare_next_tasks( # then invoke the process with the values of all non-empty channels if isinstance(proc.channels, dict): try: - val: Any = { + val: dict = { k: read_channel(channels, chan, catch=chan not in proc.triggers) for k, chan in proc.channels.items() + if isinstance(chan, str) } + + managed_values = {} + for key, chan in proc.channels.items(): + for mv in managed: + if isclass(chan) and isinstance(mv, chan): + managed_values[key] = mv( + step, PregelTaskDescription(name, val) + ) + + val.update(managed_values) except EmptyChannelError: continue elif isinstance(proc.channels, list): @@ -1464,7 +1576,12 @@ def _prepare_next_tasks( if node := proc.get_node(): tasks.append( PregelExecutableTask( - name, val, node, deque(), proc.config, triggers + name, + val, + node, + deque(), + merge_configs(config, proc.config), + triggers, ) ) else: diff --git a/langgraph/pregel/read.py b/langgraph/pregel/read.py index 6398a1e2c..e8576e22f 100644 --- a/langgraph/pregel/read.py +++ b/langgraph/pregel/read.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Any, Callable, Mapping, Optional, Sequence, Union +from typing import Any, Callable, Mapping, Optional, Sequence, Type, Union from langchain_core.pydantic_v1 import Field from langchain_core.runnables import ( @@ -15,6 +15,7 @@ from langchain_core.runnables.config import merge_configs from langchain_core.runnables.utils import ConfigurableFieldSpec from langgraph.constants import CONFIG_KEY_READ +from langgraph.managed.base import ManagedValue from langgraph.pregel.write import ChannelWrite from langgraph.utils import RunnableCallable @@ -99,7 +100,7 @@ DEFAULT_BOUND: RunnablePassthrough = RunnablePassthrough() class PregelNode(RunnableBindingBase): - channels: Union[list[str], Mapping[str, str]] + channels: Union[list[str], Mapping[str, Union[str, Type[ManagedValue]]]] triggers: list[str] = Field(default_factory=list) diff --git a/tests/__snapshots__/test_pregel.ambr b/tests/__snapshots__/test_pregel.ambr index 68e1d3a48..9e3ac9849 100644 --- a/tests/__snapshots__/test_pregel.ambr +++ b/tests/__snapshots__/test_pregel.ambr @@ -803,7 +803,7 @@ ''' # --- # name: test_prebuilt_chat - '{"title": "LangGraphInput", "$ref": "#/definitions/AgentState", "definitions": {"BaseMessage": {"title": "BaseMessage", "description": "Base abstract Message class.\\n\\nMessages are the inputs and outputs of ChatModels.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "type"]}, "AgentState": {"title": "AgentState", "type": "object", "properties": {"messages": {"title": "Messages", "type": "array", "items": {"$ref": "#/definitions/BaseMessage"}}}, "required": ["messages"]}}}' + '{"title": "LangGraphInput", "$ref": "#/definitions/AgentState", "definitions": {"BaseMessage": {"title": "BaseMessage", "description": "Base abstract Message class.\\n\\nMessages are the inputs and outputs of ChatModels.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "type"]}, "AgentState": {"title": "AgentState", "type": "object", "properties": {"messages": {"title": "Messages", "type": "array", "items": {"$ref": "#/definitions/BaseMessage"}}, "is_last_step": {"title": "Is Last Step", "type": "boolean"}}, "required": ["messages", "is_last_step"]}}}' # --- # name: test_prebuilt_chat.1 '{"title": "LangGraphOutput", "type": "object", "properties": {"messages": {"title": "Messages", "type": "array", "items": {"$ref": "#/definitions/BaseMessage"}}}, "definitions": {"BaseMessage": {"title": "BaseMessage", "description": "Base abstract Message class.\\n\\nMessages are the inputs and outputs of ChatModels.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "type"]}}}' @@ -885,7 +885,7 @@ ''' # --- # name: test_prebuilt_tool_chat - '{"title": "LangGraphInput", "$ref": "#/definitions/AgentState", "definitions": {"BaseMessage": {"title": "BaseMessage", "description": "Base abstract Message class.\\n\\nMessages are the inputs and outputs of ChatModels.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "type"]}, "AgentState": {"title": "AgentState", "type": "object", "properties": {"messages": {"title": "Messages", "type": "array", "items": {"$ref": "#/definitions/BaseMessage"}}}, "required": ["messages"]}}}' + '{"title": "LangGraphInput", "$ref": "#/definitions/AgentState", "definitions": {"BaseMessage": {"title": "BaseMessage", "description": "Base abstract Message class.\\n\\nMessages are the inputs and outputs of ChatModels.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "type"]}, "AgentState": {"title": "AgentState", "type": "object", "properties": {"messages": {"title": "Messages", "type": "array", "items": {"$ref": "#/definitions/BaseMessage"}}, "is_last_step": {"title": "Is Last Step", "type": "boolean"}}, "required": ["messages", "is_last_step"]}}}' # --- # name: test_prebuilt_tool_chat.1 '{"title": "LangGraphOutput", "type": "object", "properties": {"messages": {"title": "Messages", "type": "array", "items": {"$ref": "#/definitions/BaseMessage"}}}, "definitions": {"BaseMessage": {"title": "BaseMessage", "description": "Base abstract Message class.\\n\\nMessages are the inputs and outputs of ChatModels.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "type"]}}}' diff --git a/tests/test_pregel.py b/tests/test_pregel.py index cf3d87f12..4529e595c 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -2643,40 +2643,39 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: tools = [search_api] - app = create_tool_calling_executor( - FakeFuntionChatModel( - responses=[ - AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call123", - "name": "search_api", - "args": {"query": "query"}, - }, - ], - ), - AIMessage( - content="", - tool_calls=[ - { - "id": "tool_call234", - "name": "search_api", - "args": {"query": "another"}, - }, - { - "id": "tool_call567", - "name": "search_api", - "args": {"query": "a third one"}, - }, - ], - ), - AIMessage(content="answer"), - ] - ), - tools, + model = FakeFuntionChatModel( + responses=[ + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call123", + "name": "search_api", + "args": {"query": "query"}, + }, + ], + ), + AIMessage( + content="", + tool_calls=[ + { + "id": "tool_call234", + "name": "search_api", + "args": {"query": "another"}, + }, + { + "id": "tool_call567", + "name": "search_api", + "args": {"query": "a third one"}, + }, + ], + ), + AIMessage(content="answer"), + ] ) + app = create_tool_calling_executor(model, tools) + assert app.get_input_schema().schema_json() == snapshot assert app.get_output_schema().schema_json() == snapshot assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot @@ -2736,6 +2735,21 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None: ] } + assert app.invoke( + {"messages": [HumanMessage(content="what is weather in sf")]}, + {"recursion_limit": 2}, + debug=True, + ) == { + "messages": [ + HumanMessage(content="what is weather in sf", id=AnyStr()), + AIMessage( + content="Sorry, need more steps to process this request.", id=AnyStr() + ), + ] + } + + model.i = 0 # reset the model + assert app.invoke( {"messages": [HumanMessage(content="what is weather in sf")]}, stream_mode="updates",