From a46bac07bf78c59415249a3345b2d0fdc17c8468 Mon Sep 17 00:00:00 2001 From: Sydney Runkle Date: Wed, 6 Aug 2025 14:35:36 -0400 Subject: [PATCH] first pass --- .github/workflows/_lint.yml | 1 + .github/workflows/_test.yml | 3 + .github/workflows/ci.yml | 3 + libs/langgraph/langgraph/graph/ui.py | 5 +- libs/langgraph/langgraph/types.py | 8 +- .../langgraph/prebuilt/chat_agent_executor.py | 124 +++++++++----- libs/prebuilt/langgraph/prebuilt/tool_node.py | 153 +++++++++++++----- .../langgraph/prebuilt/tool_validator.py | 76 ++++++--- testing.py | 67 ++++++++ 9 files changed, 330 insertions(+), 110 deletions(-) create mode 100644 testing.py diff --git a/.github/workflows/_lint.yml b/.github/workflows/_lint.yml index 68c9a78f3..e88c54364 100644 --- a/.github/workflows/_lint.yml +++ b/.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: diff --git a/.github/workflows/_test.yml b/.github/workflows/_test.yml index 6fb763003..2a53afc12 100644 --- a/.github/workflows/_test.yml +++ b/.github/workflows/_test.yml @@ -11,6 +11,9 @@ on: permissions: contents: read +env: + UV_PRERELEASE: allow + jobs: build: runs-on: ubuntu-latest diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 19a735069..3b0e5b33c 100644 --- a/.github/workflows/ci.yml +++ b/.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. # diff --git a/libs/langgraph/langgraph/graph/ui.py b/libs/langgraph/langgraph/graph/ui.py index f2fe5a1c2..04d11e225 100644 --- a/libs/langgraph/langgraph/graph/ui.py +++ b/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: diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 5fbc304b7..0ba5c7e20 100644 --- a/libs/langgraph/langgraph/types.py +++ b/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", diff --git a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py index a68213702..4026ef81f 100644 --- a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py +++ b/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 ':' 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 diff --git a/libs/prebuilt/langgraph/prebuilt/tool_node.py b/libs/prebuilt/langgraph/prebuilt/tool_node.py index 252f88802..3bf84243a 100644 --- a/libs/prebuilt/langgraph/prebuilt/tool_node.py +++ b/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__" diff --git a/libs/prebuilt/langgraph/prebuilt/tool_validator.py b/libs/prebuilt/langgraph/prebuilt/tool_validator.py index 19df4233f..3b88d5e7b 100644 --- a/libs/prebuilt/langgraph/prebuilt/tool_validator.py +++ b/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)] diff --git a/testing.py b/testing.py new file mode 100644 index 000000000..bb02b053c --- /dev/null +++ b/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") \ No newline at end of file