first pass

.github/workflows/_lint.yml

@@ -14,6 +14,7 @@ permissions:
 env:
   # This env var allows us to get inline annotations when ruff has complaints.
   RUFF_OUTPUT_FORMAT: github
+  UV_PRERELEASE: allow
 
 jobs:
   build:

.github/workflows/_test.yml

@@ -11,6 +11,9 @@ on:
 permissions:
   contents: read
 
+env:
+  UV_PRERELEASE: allow
+
 jobs:
   build:
     runs-on: ubuntu-latest

.github/workflows/ci.yml

@@ -10,6 +10,9 @@ on:
 permissions:
   contents: read
 
+env:
+  UV_PRERELEASE: allow
+
 # If another push to the same PR or branch happens while this workflow is still running,
 # cancel the earlier run in favor of the next run.
 #

libs/langgraph/langgraph/graph/ui.py

@@ -3,7 +3,8 @@ from __future__ import annotations
 from typing import Any, Literal, Union, cast
 from uuid import uuid4
 
-from langchain_core.messages import AnyMessage
+from langchain_core.messages import AnyMessage as MessageV0
+from langchain_core.v1.messages import MessageV1 as MessageV1
 from typing_extensions import TypedDict
 
 from langgraph.config import get_config, get_stream_writer
@@ -64,7 +65,7 @@ def push_ui_message(
     *,
     id: str | None = None,
     metadata: dict[str, Any] | None = None,
-    message: AnyMessage | None = None,
+    message: MessageV0 | MessageV1 | None = None,
     state_key: str | None = "ui",
     merge: bool = False,
 ) -> UIMessage:

libs/langgraph/langgraph/types.py

@@ -33,13 +33,7 @@ if TYPE_CHECKING:
     from langgraph.pregel.protocol import PregelProtocol
 
 
-try:
-    from langchain_core.messages.tool import ToolOutputMixin
-except ImportError:
-
-    class ToolOutputMixin:  # type: ignore[no-redef]
-        pass
-
+from langchain_core.messages.tool import ToolOutputMixin
 
 __all__ = (
     "All",

libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py

@@ -15,22 +15,26 @@ from typing import (
 from warnings import warn
 
 from langchain_core.language_models import (
-    BaseChatModel,
+    BaseChatModel as BaseChatModelV0,
+)
+from langchain_core.language_models import (
     LanguageModelInput,
     LanguageModelLike,
 )
 from langchain_core.messages import (
     AIMessage as AIMessageV0,
-    AnyMessage as MessageV0,
-    BaseMessage as BaseMessageV0,
-    SystemMessage as SystemMessageV0,
-    ToolMessage as ToolMessageV0,
 )
-from langchain_core.v1.messages import(
-    AIMessage as AIMessageV1,
-    MessageV1,
-    SystemMessage as SystemMessageV1,
-    ToolMessage as ToolMessageV1,
+from langchain_core.messages import (
+    AnyMessage as MessageV0,
+)
+from langchain_core.messages import (
+    BaseMessage as BaseMessageV0,
+)
+from langchain_core.messages import (
+    SystemMessage as SystemMessageV0,
+)
+from langchain_core.messages import (
+    ToolMessage as ToolMessageV0,
 )
 from langchain_core.runnables import (
     Runnable,
@@ -39,6 +43,19 @@ from langchain_core.runnables import (
     RunnableSequence,
 )
 from langchain_core.tools import BaseTool
+from langchain_core.v1.chat_models import BaseChatModel as BaseChatModelV1
+from langchain_core.v1.messages import (
+    AIMessage as AIMessageV1,
+)
+from langchain_core.v1.messages import (
+    MessageV1,
+)
+from langchain_core.v1.messages import (
+    SystemMessage as SystemMessageV1,
+)
+from langchain_core.v1.messages import (
+    ToolMessage as ToolMessageV1,
+)
 from pydantic import BaseModel
 from typing_extensions import Annotated, NotRequired, TypedDict
 
@@ -115,13 +132,16 @@ def _get_state_value(state: StateSchema, key: str, default: Any = None) -> Any:
     )
 
 
-def _get_prompt_runnable(prompt: Optional[Prompt], message_version: Literal["v0", "v1"]) -> Runnable:
+def _get_prompt_runnable(
+    prompt: Optional[Prompt], message_version: Literal["v0", "v1"]
+) -> Runnable:
     prompt_runnable: Runnable
     if prompt is None:
         prompt_runnable = RunnableCallable(
             lambda state: _get_state_value(state, "messages"), name=PROMPT_RUNNABLE_NAME
         )
     elif isinstance(prompt, str):
+        _system_message: Union[SystemMessageV0, SystemMessageV1]
         if message_version == "v0":
             _system_message = SystemMessageV0(content=prompt)
         else:
@@ -162,7 +182,7 @@ def _should_bind_tools(
             (
                 step
                 for step in model.steps
-                if isinstance(step, (RunnableBinding, BaseChatModel))
+                if isinstance(step, (RunnableBinding, BaseChatModelV0))
             ),
             model,
         )
@@ -201,14 +221,14 @@ def _should_bind_tools(
     return False
 
 
-def _get_model(model: LanguageModelLike) -> BaseChatModel:
+def _get_model(model: LanguageModelLike) -> Union[BaseChatModelV0, BaseChatModelV1]:
     """Get the underlying model from a RunnableBinding or return the model itself."""
     if isinstance(model, RunnableSequence):
         model = next(
             (
                 step
                 for step in model.steps
-                if isinstance(step, (RunnableBinding, BaseChatModel))
+                if isinstance(step, (RunnableBinding, BaseChatModelV0))
             ),
             model,
         )
@@ -216,7 +236,7 @@ def _get_model(model: LanguageModelLike) -> BaseChatModel:
     if isinstance(model, RunnableBinding):
         model = model.bound
 
-    if not isinstance(model, BaseChatModel):
+    if not isinstance(model, (BaseChatModelV0, BaseChatModelV1)):
         raise TypeError(
             f"Expected `model` to be a ChatModel or RunnableBinding (e.g. model.bind_tools(...)), got {type(model)}"
         )
@@ -235,7 +255,9 @@ def _validate_chat_history(
         for tool_call in message.tool_calls
     ]
     tool_call_ids_with_results = {
-        message.tool_call_id for message in messages if isinstance(message, (ToolMessageV0, ToolMessageV1))
+        message.tool_call_id
+        for message in messages
+        if isinstance(message, (ToolMessageV0, ToolMessageV1))
     }
     tool_calls_without_results = [
         tool_call
@@ -259,10 +281,16 @@ def create_react_agent(
     model: Union[
         str,
         LanguageModelLike,
-        Callable[[StateSchema, Runtime[ContextT]], BaseChatModel],
-        Callable[[StateSchema, Runtime[ContextT]], Awaitable[BaseChatModel]],
         Callable[
-            [StateSchema, Runtime[ContextT]], Runnable[LanguageModelInput, Union[BaseMessageV0, MessageV1]]
+            [StateSchema, Runtime[ContextT]], Union[BaseChatModelV0, BaseChatModelV1]
+        ],
+        Callable[
+            [StateSchema, Runtime[ContextT]],
+            Awaitable[Union[BaseChatModelV0, BaseChatModelV1]],
+        ],
+        Callable[
+            [StateSchema, Runtime[ContextT]],
+            Runnable[LanguageModelInput, Union[BaseMessageV0, MessageV1]],
         ],
         Callable[
             [StateSchema, Runtime[ContextT]],
@@ -425,6 +453,7 @@ def create_react_agent(
         name: An optional name for the CompiledStateGraph.
             This name will be automatically used when adding ReAct agent graph to another graph as a subgraph node -
             particularly useful for building multi-agent systems.
+        message_version: Version of the message format to use (v0 or v1 langchain).
 
     !!! warning "`config_schema` Deprecated"
         The `config_schema` parameter is deprecated in v0.6.0 and support will be removed in v2.0.0.
@@ -488,7 +517,8 @@ def create_react_agent(
         raise ValueError(
             f"Invalid version {version}. Supported versions are 'v1' and 'v2'."
         )
-    
+
+    message_constructor: Union[type[AIMessageV0], type[AIMessageV1]]
     if message_version == "v0":
         message_constructor = AIMessageV0
     else:
@@ -516,7 +546,10 @@ def create_react_agent(
         tool_node = tools
     else:
         llm_builtin_tools = [t for t in tools if isinstance(t, dict)]
-        tool_node = ToolNode([t for t in tools if not isinstance(t, dict)])
+        tool_node = ToolNode(
+            [t for t in tools if not isinstance(t, dict)],
+            message_version=message_version,
+        )
         tool_classes = list(tool_node.tools_by_name.values())
 
     is_dynamic_model = not isinstance(model, (str, Runnable)) and callable(model)
@@ -536,17 +569,19 @@ def create_react_agent(
                     "use '<provider>:<model>' string syntax for `model` parameter."
                 )
 
-            model = cast(BaseChatModel, init_chat_model(model))
+            model = init_chat_model(model, message_version=message_version)
 
         if (
             _should_bind_tools(model, tool_classes, num_builtin=len(llm_builtin_tools))  # type: ignore[arg-type]
             and len(tool_classes + llm_builtin_tools) > 0
         ):
-            model = cast(BaseChatModel, model).bind_tools(
+            model = cast(Union[BaseChatModelV0, BaseChatModelV1], model).bind_tools(
                 tool_classes + llm_builtin_tools  # type: ignore[operator]
             )
 
-        static_model: Optional[Runnable] = _get_prompt_runnable(prompt, message_version) | model  # type: ignore[operator]
+        static_model: Optional[Runnable] = (
+            _get_prompt_runnable(prompt, message_version) | model  # type: ignore[operator]
+        )
     else:
         # For dynamic models, we'll create the runnable at runtime
         static_model = None
@@ -560,7 +595,7 @@ def create_react_agent(
     ) -> LanguageModelLike:
         """Resolve the model to use, handling both static and dynamic models."""
         if is_dynamic_model:
-            return _get_prompt_runnable(prompt) | model(state, runtime)  # type: ignore[operator]
+            return _get_prompt_runnable(prompt, message_version) | model(state, runtime)  # type: ignore[operator]
         else:
             return static_model
 
@@ -570,14 +605,18 @@ def create_react_agent(
         """Async resolve the model to use, handling both static and dynamic models."""
         if is_async_dynamic_model:
             resolved_model = await model(state, runtime)  # type: ignore[misc,operator]
-            return _get_prompt_runnable(prompt) | resolved_model
+            return _get_prompt_runnable(prompt, message_version) | resolved_model
         elif is_dynamic_model:
-            return _get_prompt_runnable(prompt) | model(state, runtime)  # type: ignore[operator]
+            return _get_prompt_runnable(prompt, message_version) | model(state, runtime)  # type: ignore[operator]
         else:
             return static_model
 
-    def _are_more_steps_needed(state: StateSchema, response: Union[BaseMessageV0, MessageV1]) -> bool:
-        has_tool_calls = isinstance(response, (AIMessageV0, AIMessageV1)) and response.tool_calls
+    def _are_more_steps_needed(
+        state: StateSchema, response: Union[AIMessageV0, AIMessageV1]
+    ) -> bool:
+        has_tool_calls = (
+            isinstance(response, (AIMessageV0, AIMessageV1)) and response.tool_calls
+        )
         all_tools_return_direct = (
             all(call["name"] in should_return_direct for call in response.tool_calls)
             if isinstance(response, (AIMessageV0, AIMessageV1))
@@ -691,7 +730,6 @@ def create_react_agent(
 
     input_schema: StateSchemaType
     if pre_model_hook is not None:
-
         # Dynamically create a schema that inherits from state_schema and adds 'llm_input_messages'
         if isinstance(state_schema, type) and issubclass(state_schema, BaseModel):
             # For Pydantic schemas
@@ -699,10 +737,10 @@ def create_react_agent(
 
             if message_version == "v0":
                 input_schema = create_model(
-                "CallModelInputSchema",
-                llm_input_messages=(list[MessageV0], ...),
-                __base__=state_schema,
-            )
+                    "CallModelInputSchema",
+                    llm_input_messages=(list[MessageV0], ...),
+                    __base__=state_schema,
+                )
             else:
                 input_schema = create_model(
                     "CallModelInputSchema",
@@ -710,11 +748,12 @@ def create_react_agent(
                     __base__=state_schema,
                 )
         else:
-
             if message_version == "v0":
+
                 class CallModelInputSchema(state_schema):  # type: ignore
                     llm_input_messages: list[MessageV0]
             else:
+
                 class CallModelInputSchema(state_schema):  # type: ignore
                     llm_input_messages: list[MessageV1]
 
@@ -819,7 +858,10 @@ def create_react_agent(
         messages = _get_state_value(state, "messages")
         last_message = messages[-1]
         # If there is no function call, then we finish
-        if not isinstance(last_message, (AIMessageV0, AIMessageV1)) or not last_message.tool_calls:
+        if (
+            not isinstance(last_message, (AIMessageV0, AIMessageV1))
+            or not last_message.tool_calls
+        ):
             if post_model_hook is not None:
                 return "post_model_hook"
             elif response_format is not None:
@@ -908,10 +950,14 @@ def create_react_agent(
 
             messages = _get_state_value(state, "messages")
             tool_messages = [
-                m.tool_call_id for m in messages if isinstance(m, (ToolMessageV0, ToolMessageV1)):
+                m.tool_call_id
+                for m in messages
+                if isinstance(m, (ToolMessageV0, ToolMessageV1))
             ]
             last_ai_message = next(
-                m for m in reversed(messages) if isinstance(m, (AIMessageV0, AIMessageV1))
+                m
+                for m in reversed(messages)
+                if isinstance(m, (AIMessageV0, AIMessageV1))
             )
             pending_tool_calls = [
                 c for c in last_ai_message.tool_calls if c["id"] not in tool_messages
@@ -944,7 +990,7 @@ def create_react_agent(
 
     def route_tool_responses(state: StateSchema) -> str:
         for m in reversed(_get_state_value(state, "messages")):
-            if not isinstance(m, ToolMessage):
+            if not isinstance(m, (ToolMessageV0, ToolMessageV1)):
                 break
             if m.name in should_return_direct:
                 return END

libs/prebuilt/langgraph/prebuilt/tool_node.py

@@ -50,12 +50,26 @@ from typing import (
 )
 
 from langchain_core.messages import (
-    AIMessage,
-    AnyMessage,
-    RemoveMessage,
-    ToolCall,
-    ToolMessage,
-    convert_to_messages,
+    AIMessage as AIMessageV0,
+)
+from langchain_core.messages import (
+    AIMessage as AIMessageV1,
+)
+from langchain_core.messages import (
+    AnyMessage as MessageV0,
+)
+from langchain_core.messages import (
+    RemoveMessage as RemoveMessageV0,
+)
+from langchain_core.messages import (
+    ToolMessage as ToolMessageV0,
+)
+from langchain_core.messages.tool import ToolCall
+from langchain_core.messages.utils import (
+    convert_to_messages as convert_to_messages_v0,
+)
+from langchain_core.messages.utils import (
+    convert_to_messages_v1,
 )
 from langchain_core.runnables import RunnableConfig
 from langchain_core.runnables.config import (
@@ -68,6 +82,12 @@ from langchain_core.tools.base import (
     TOOL_MESSAGE_BLOCK_TYPES,
     get_all_basemodel_annotations,
 )
+from langchain_core.v1.messages import (
+    MessageV1,
+)
+from langchain_core.v1.messages import (
+    ToolMessage as ToolMessageV1,
+)
 from pydantic import BaseModel
 from typing_extensions import Annotated, get_args, get_origin
 
@@ -326,6 +346,7 @@ class ToolNode(RunnableCallable):
             bool, str, Callable[..., str], tuple[type[Exception], ...]
         ] = True,
         messages_key: str = "messages",
+        message_version: Literal["v0", "v1"] = "v0",
     ) -> None:
         """Initialize the ToolNode with the provided tools and configuration.
 
@@ -342,6 +363,7 @@ class ToolNode(RunnableCallable):
         self.tool_to_store_arg: dict[str, Optional[str]] = {}
         self.handle_tool_errors = handle_tool_errors
         self.messages_key = messages_key
+        self.message_version = message_version
         for tool_ in tools:
             if not isinstance(tool_, BaseTool):
                 tool_ = create_tool(tool_)
@@ -352,7 +374,7 @@ class ToolNode(RunnableCallable):
     def _func(
         self,
         input: Union[
-            list[AnyMessage],
+            list[Union[MessageV0, MessageV1]],
             dict[str, Any],
             BaseModel,
         ],
@@ -373,7 +395,7 @@ class ToolNode(RunnableCallable):
     async def _afunc(
         self,
         input: Union[
-            list[AnyMessage],
+            list[Union[MessageV0, MessageV1]],
             dict[str, Any],
             BaseModel,
         ],
@@ -390,9 +412,15 @@ class ToolNode(RunnableCallable):
 
     def _combine_tool_outputs(
         self,
-        outputs: list[ToolMessage],
+        outputs: list[Union[ToolMessageV0, ToolMessageV1, Command]],
         input_type: Literal["list", "dict", "tool_calls"],
-    ) -> list[Union[Command, list[ToolMessage], dict[str, list[ToolMessage]]]]:
+    ) -> list[
+        Union[
+            Command,
+            list[Union[ToolMessageV0, ToolMessageV1]],
+            dict[str, list[Union[ToolMessageV0, ToolMessageV1]]],
+        ]
+    ]:
         # preserve existing behavior for non-command tool outputs for backwards
         # compatibility
         if not any(isinstance(output, Command) for output in outputs):
@@ -402,7 +430,9 @@ class ToolNode(RunnableCallable):
         # LangGraph will automatically handle list of Command and non-command node
         # updates
         combined_outputs: list[
-            Command | list[ToolMessage] | dict[str, list[ToolMessage]]
+            Command
+            | list[Union[ToolMessageV0, ToolMessageV1]]
+            | dict[str, list[Union[ToolMessageV0, ToolMessageV1]]]
         ] = []
 
         # combine all parent commands with goto into a single parent command
@@ -437,7 +467,7 @@ class ToolNode(RunnableCallable):
         call: ToolCall,
         input_type: Literal["list", "dict", "tool_calls"],
         config: RunnableConfig,
-    ) -> ToolMessage:
+    ) -> Union[ToolMessageV0, ToolMessageV1, Command]:
         """Run a single tool call synchronously."""
         if invalid_tool_message := self._validate_tool_call(call):
             return invalid_tool_message
@@ -469,20 +499,31 @@ class ToolNode(RunnableCallable):
             # Handled
             else:
                 content = _handle_tool_error(e, flag=self.handle_tool_errors)
-            return ToolMessage(
-                content=content,
-                name=call["name"],
-                tool_call_id=call["id"],
-                status="error",
-            )
+            if self.message_version == "v0":
+                return ToolMessageV0(
+                    content=content,
+                    name=call["name"],
+                    tool_call_id=call["id"],
+                    status="error",
+                )
+            else:
+                return ToolMessageV1(
+                    content=content,
+                    name=call["name"],
+                    tool_call_id=cast(str, call["id"]),
+                    status="error",
+                )
 
         if isinstance(response, Command):
             return self._validate_tool_command(response, call, input_type)
-        elif isinstance(response, ToolMessage):
+        elif isinstance(response, ToolMessageV0):
             response.content = cast(
                 Union[str, list], msg_content_output(response.content)
             )
             return response
+        elif isinstance(response, ToolMessageV1):
+            # TODO: Handle ToolMessageV1
+            return response
         else:
             raise TypeError(
                 f"Tool {call['name']} returned unexpected type: {type(response)}"
@@ -493,7 +534,7 @@ class ToolNode(RunnableCallable):
         call: ToolCall,
         input_type: Literal["list", "dict", "tool_calls"],
         config: RunnableConfig,
-    ) -> ToolMessage:
+    ) -> Union[ToolMessageV0, ToolMessageV1, Command]:
         """Run a single tool call asynchronously."""
         if invalid_tool_message := self._validate_tool_call(call):
             return invalid_tool_message
@@ -527,20 +568,31 @@ class ToolNode(RunnableCallable):
             else:
                 content = _handle_tool_error(e, flag=self.handle_tool_errors)
 
-            return ToolMessage(
-                content=content,
-                name=call["name"],
-                tool_call_id=call["id"],
-                status="error",
-            )
+            if self.message_version == "v0":
+                return ToolMessageV0(
+                    content=content,
+                    name=call["name"],
+                    tool_call_id=call["id"],
+                    status="error",
+                )
+            else:
+                return ToolMessageV1(
+                    content=content,
+                    name=call["name"],
+                    tool_call_id=cast(str, call["id"]),
+                    status="error",
+                )
 
         if isinstance(response, Command):
             return self._validate_tool_command(response, call, input_type)
-        elif isinstance(response, ToolMessage):
+        elif isinstance(response, ToolMessageV0):
             response.content = cast(
                 Union[str, list], msg_content_output(response.content)
             )
             return response
+        elif isinstance(response, ToolMessageV1):
+            # TODO: Handle ToolMessageV1
+            return response
         else:
             raise TypeError(
                 f"Tool {call['name']} returned unexpected type: {type(response)}"
@@ -549,7 +601,7 @@ class ToolNode(RunnableCallable):
     def _parse_input(
         self,
         input: Union[
-            list[AnyMessage],
+            list[Union[MessageV0, MessageV1]],
             dict[str, Any],
             BaseModel,
         ],
@@ -574,7 +626,9 @@ class ToolNode(RunnableCallable):
 
         try:
             latest_ai_message = next(
-                m for m in reversed(messages) if isinstance(m, AIMessage)
+                m
+                for m in reversed(messages)
+                if isinstance(m, (AIMessageV0, AIMessageV1))
             )
         except StopIteration:
             raise ValueError("No AIMessage found in input")
@@ -585,15 +639,29 @@ class ToolNode(RunnableCallable):
         ]
         return tool_calls, input_type
 
-    def _validate_tool_call(self, call: ToolCall) -> Optional[ToolMessage]:
+    def _validate_tool_call(
+        self, call: ToolCall
+    ) -> Optional[Union[ToolMessageV0, ToolMessageV1]]:
         if (requested_tool := call["name"]) not in self.tools_by_name:
             content = INVALID_TOOL_NAME_ERROR_TEMPLATE.format(
                 requested_tool=requested_tool,
                 available_tools=", ".join(self.tools_by_name.keys()),
             )
-            return ToolMessage(
-                content, name=requested_tool, tool_call_id=call["id"], status="error"
-            )
+
+            if self.message_version == "v0":
+                return ToolMessageV0(
+                    content=content,
+                    name=requested_tool,
+                    tool_call_id=call["id"],
+                    status="error",
+                )
+            else:
+                return ToolMessageV1(
+                    content=content,
+                    name=requested_tool,
+                    tool_call_id=cast(str, call["id"]),
+                    status="error",
+                )
         else:
             return None
 
@@ -601,7 +669,7 @@ class ToolNode(RunnableCallable):
         self,
         tool_call: ToolCall,
         input: Union[
-            list[AnyMessage],
+            list[Union[MessageV0, MessageV1]],
             dict[str, Any],
             BaseModel,
         ],
@@ -665,7 +733,7 @@ class ToolNode(RunnableCallable):
         self,
         tool_call: ToolCall,
         input: Union[
-            list[AnyMessage],
+            list[Union[MessageV0, MessageV1]],
             dict[str, Any],
             BaseModel,
         ],
@@ -738,20 +806,23 @@ class ToolNode(RunnableCallable):
                 )
 
             updated_command = deepcopy(command)
-            messages_update = updated_command.update
+            messages_update = updated_command.update or []
         else:
             return command
 
         # convert to message objects if updates are in a dict format
-        messages_update = convert_to_messages(messages_update)
+        if self.message_version == "v0":
+            messages_update = convert_to_messages_v0(messages_update)
+        else:
+            messages_update = convert_to_messages_v1(messages_update)
 
         # no validation needed if all messages are being removed
-        if messages_update == [RemoveMessage(id=REMOVE_ALL_MESSAGES)]:
+        if messages_update == [RemoveMessageV0(id=REMOVE_ALL_MESSAGES)]:
             return updated_command
 
         has_matching_tool_message = False
         for message in messages_update:
-            if not isinstance(message, ToolMessage):
+            if not isinstance(message, (ToolMessageV0, ToolMessageV1)):
                 continue
 
             if message.tool_call_id == call["id"]:
@@ -775,7 +846,7 @@ class ToolNode(RunnableCallable):
 
 
 def tools_condition(
-    state: Union[list[AnyMessage], dict[str, Any], BaseModel],
+    state: Union[list[Union[MessageV0, MessageV1]], dict[str, Any], BaseModel],
     messages_key: str = "messages",
 ) -> Literal["tools", "__end__"]:
     """Conditional routing function for tool-calling workflows.
@@ -846,7 +917,7 @@ def tools_condition(
         ai_message = messages[-1]
     else:
         raise ValueError(f"No messages found in input state to tool_edge: {state}")
-    if hasattr(ai_message, "tool_calls") and len(ai_message.tool_calls) > 0:
+    if (tool_calls := getattr(ai_message, "tool_calls", [])) and len(tool_calls) > 0:
         return "tools"
     return "__end__"
 

libs/prebuilt/langgraph/prebuilt/tool_validator.py

@@ -9,6 +9,7 @@ from typing import (
     Any,
     Callable,
     Dict,
+    Literal,
     Optional,
     Sequence,
     Tuple,
@@ -18,17 +19,30 @@ from typing import (
 )
 
 from langchain_core.messages import (
-    AIMessage,
-    AnyMessage,
-    ToolCall,
-    ToolMessage,
+    AIMessage as AIMessageV0,
 )
+from langchain_core.messages import (
+    AnyMessage as MessageV0,
+)
+from langchain_core.messages import (
+    ToolMessage as ToolMessageV0,
+)
+from langchain_core.messages.tool import ToolCall
 from langchain_core.runnables import (
     RunnableConfig,
 )
 from langchain_core.runnables.config import get_executor_for_config
 from langchain_core.tools import BaseTool, create_schema_from_function
 from langchain_core.utils.pydantic import is_basemodel_subclass
+from langchain_core.v1.messages import (
+    AIMessage as AIMessageV1,
+)
+from langchain_core.v1.messages import (
+    MessageV1,
+)
+from langchain_core.v1.messages import (
+    ToolMessage as ToolMessageV1,
+)
 from pydantic import BaseModel, ValidationError
 from pydantic.v1 import BaseModel as BaseModelV1
 from pydantic.v1 import ValidationError as ValidationErrorV1
@@ -68,6 +82,7 @@ class ValidationNode(RunnableCallable):
             exception repr and a message to respond after fixing validation errors.
         name: The name of the node.
         tags: A list of tags to add to the node.
+        message_version: Version of the message format to use (v0 or v1 langchain).
 
     Returns:
         (Union[Dict[str, List[ToolMessage]], Sequence[ToolMessage]]): A list of ToolMessages with the validated content or error messages.
@@ -134,6 +149,7 @@ class ValidationNode(RunnableCallable):
         ] = None,
         name: str = "validation",
         tags: Optional[list[str]] = None,
+        message_version: Literal["v0", "v1"] = "v0",
     ) -> None:
         super().__init__(self._func, None, name=name, tags=tags, trace=False)
         self._format_error = format_error or _default_format_error
@@ -163,10 +179,11 @@ class ValidationNode(RunnableCallable):
                 raise ValueError(
                     f"Unsupported input to ValidationNode. Expected BaseModel, tool or function. Got: {type(schema)}."
                 )
+        self.message_version = message_version
 
     def _get_message(
-        self, input: Union[list[AnyMessage], dict[str, Any]]
-    ) -> Tuple[str, AIMessage]:
+        self, input: Union[list[Union[MessageV0, MessageV1]], dict[str, Any]]
+    ) -> Tuple[str, Union[AIMessageV0, AIMessageV1]]:
         """Extract the last AIMessage from the input."""
         if isinstance(input, list):
             output_type = "list"
@@ -175,18 +192,19 @@ class ValidationNode(RunnableCallable):
             output_type = "dict"
         else:
             raise ValueError("No message found in input")
-        message: AnyMessage = messages[-1]
-        if not isinstance(message, AIMessage):
+        if not isinstance((message := messages[-1]), (AIMessageV0, AIMessageV1)):
             raise ValueError("Last message is not an AIMessage")
         return output_type, message
 
     def _func(
-        self, input: Union[list[AnyMessage], dict[str, Any]], config: RunnableConfig
+        self,
+        input: Union[list[Union[MessageV0, MessageV1]], dict[str, Any]],
+        config: RunnableConfig,
     ) -> Any:
         """Validate and run tool calls synchronously."""
         output_type, message = self._get_message(input)
 
-        def run_one(call: ToolCall) -> ToolMessage:
+        def run_one(call: ToolCall) -> Union[ToolMessageV0, ToolMessageV1]:
             schema = self.schemas_by_name[call["name"]]
             try:
                 if issubclass(schema, BaseModel):
@@ -199,18 +217,34 @@ class ValidationNode(RunnableCallable):
                     raise ValueError(
                         f"Unsupported schema type: {type(schema)}. Expected BaseModel or BaseModelV1."
                     )
-                return ToolMessage(
-                    content=content,
-                    name=call["name"],
-                    tool_call_id=cast(str, call["id"]),
-                )
+
+                if self.message_version == "v0":
+                    return ToolMessageV0(
+                        content=content,
+                        name=call["name"],
+                        tool_call_id=cast(str, call["id"]),
+                        additional_kwargs={"is_error": False},
+                    )
+                else:
+                    return ToolMessageV1(
+                        content=content,
+                        name=call["name"],
+                        tool_call_id=cast(str, call["id"]),
+                    )
             except (ValidationError, ValidationErrorV1) as e:
-                return ToolMessage(
-                    content=self._format_error(e, call, schema),
-                    name=call["name"],
-                    tool_call_id=cast(str, call["id"]),
-                    additional_kwargs={"is_error": True},
-                )
+                if self.message_version == "v0":
+                    return ToolMessageV0(
+                        content=self._format_error(e, call, schema),
+                        name=call["name"],
+                        tool_call_id=cast(str, call["id"]),
+                        additional_kwargs={"is_error": True},
+                    )
+                else:
+                    return ToolMessageV1(
+                        content=self._format_error(e, call, schema),
+                        name=call["name"],
+                        tool_call_id=cast(str, call["id"]),
+                    )
 
         with get_executor_for_config(config) as executor:
             outputs = [*executor.map(run_one, message.tool_calls)]

testing.py

@@ -0,0 +1,67 @@
+# import asyncio
+# from langchain_openai import ChatOpenAI
+# from langgraph.prebuilt import create_react_agent
+# from langchain_core.tools import tool
+# from langgraph.graph import END, START, StateGraph
+# from langgraph.graph import MessagesState
+
+
+# @tool
+# def get_weather(city: str) -> str:
+#     """
+#     Get the weather of a city
+#     """
+#     return f"The weather of {city} is sunny."
+
+
+# agent = create_react_agent(
+#     model=ChatOpenAI(
+#         model="gpt-4.1-mini",
+#         temperature=0,
+#     ),
+#     prompt="""
+#             You are a helpful travel assistant that can help user to get travel information.
+#             When providing travel information, please also include:
+#             1. Top tourist attractions and landmarks
+#             2. Any travel recommendations based on the city weather
+#             """,
+#     tools=[get_weather],
+# )
+
+
+# async def node(state: MessagesState) -> MessagesState:
+#     print("BEGIN")
+#     msg_content = ""
+
+#     async for ns, msg in agent.astream(
+#         {
+#             "messages": [
+#                 ("user", state["messages"][-1].content),
+#             ]
+#         },
+#         stream_mode="messages",
+#         # subgraphs=True,
+
+#     ):
+#         msg_content += msg[0].content
+#     print("END")
+
+#     return {"messages": [("assistant", msg_content)]}
+
+
+# graph = StateGraph(state_schema=MessagesState)
+# graph.add_node("node", node)
+# graph.add_edge(START, "node")
+# graph.add_edge("node", END)
+# workflow = graph.compile()
+
+# async def main():
+#     result = await workflow.ainvoke({"messages": [("user", "What is the weather in Tokyo?")]})
+#     print(result)
+
+# if __name__ == "__main__":
+#     asyncio.run(main())
+
+from langchain.chat_models import init_chat_model
+
+model = init_chat_model("openai:gpt-4o-mini", message_version="v1")