WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-11 12:17:53 +02:00
Code Issues Packages Projects Releases Wiki Activity

Add managed values and IsLastStep (#330)

Browse Source 
* 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 <andrewnguonly@gmail.com>
This commit is contained in:
Nuno Campos
2024-05-10 13:44:20 -07:00
committed by GitHub
co-authored by Andrew Nguonly
parent b33aed6f67
commit 445a110917
9 changed files with 369 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files
langgraph/graph/state.py
+36 -10
View File
@@ -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
langgraph/managed/__init__.py
+3
View File
@@ -0,0 +1,3 @@
from langgraph.managed.is_last_step import IsLastStep
__all__ = ["IsLastStep"]
langgraph/managed/base.py
+105
View File
@@ -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
)
)
langgraph/managed/is_last_step.py
+12
View File
@@ -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]
langgraph/prebuilt/chat_agent_executor.py
+27 -1
View File
@@ -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]}
langgraph/pregel/__init__.py
+135 -18
View File
@@ -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:
langgraph/pregel/read.py
+3 -2
View File
@@ -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)
tests/__snapshots__/test_pregel.ambr
+2 -2
View File
@@ -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"]}}}'
tests/test_pregel.py
+46 -32
View File
@@ -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",