diff --git a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py index eecc37837..749d81807 100644 --- a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py +++ b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py @@ -564,6 +564,102 @@ class _AgentBuilder: return generate_structured_response, agenerate_structured_response + def _create_model_router(self, entrypoint: str) -> Callable: + """Create the model router function (should_continue).""" + def should_continue(state: StateSchema) -> Union[str, list[Send]]: + messages = _get_state_value(state, "messages") + last_message = messages[-1] + # If there is no function call, then we finish + if not isinstance(last_message, AIMessage) or not last_message.tool_calls: + if self.post_model_hook is not None: + return "post_model_hook" + elif self.response_format is not None: + return "generate_structured_response" + else: + return END + # Otherwise if there is, we continue + else: + if self.version == "v1": + return "tools" + elif self.version == "v2": + if self.post_model_hook is not None: + return "post_model_hook" + return [ + Send( + "tools", + ToolCallWithContext( + __type="tool_call_with_context", + tool_call=tool_call, + state=state, + ), + ) + for tool_call in last_message.tool_calls + ] + + return should_continue + + def _create_tools_router(self, entrypoint: str) -> Callable: + """Create the tools router function (route_tool_responses).""" + def route_tool_responses(state: StateSchema) -> str: + for m in reversed(_get_state_value(state, "messages")): + if not isinstance(m, ToolMessage): + break + if m.name in self.should_return_direct: + return END + + # handle a case of parallel tool calls where + # the tool w/ `return_direct` was executed in a different `Send` + if isinstance(m, AIMessage) and m.tool_calls: + if any(call["name"] in self.should_return_direct for call in m.tool_calls): + return END + + return entrypoint + + return route_tool_responses + + def _create_post_model_hook_router(self, entrypoint: str) -> Callable: + """Create the post model hook router function.""" + def post_model_hook_router(state: StateSchema) -> Union[str, list[Send]]: + """Route to the next node after post_model_hook. + + Routes to one of: + * "tools": if there are pending tool calls without a corresponding message. + * "generate_structured_response": if no pending tool calls exist and response_format is specified. + * END: if no pending tool calls exist and no response_format is specified. + """ + + messages = _get_state_value(state, "messages") + tool_messages = [ + m.tool_call_id for m in messages if isinstance(m, ToolMessage) + ] + last_ai_message = next( + m for m in reversed(messages) if isinstance(m, AIMessage) + ) + pending_tool_calls = [ + c for c in last_ai_message.tool_calls if c["id"] not in tool_messages + ] + + if pending_tool_calls: + return [ + Send( + "tools", + ToolCallWithContext( + __type="tool_call_with_context", + tool_call=tool_call, + state=state, + ), + ) + for tool_call in pending_tool_calls + ] + elif isinstance(messages[-1], ToolMessage): + return entrypoint + elif self.response_format is not None: + return "generate_structured_response" + else: + return END + + return post_model_hook_router + def create_react_agent( model: Union[ @@ -1273,3 +1369,4 @@ __all__ = [ +