diff --git a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py index b52fc9580..26f0ce8d9 100644 --- a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py +++ b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py @@ -698,8 +698,63 @@ class _AgentBuilder: def _setup_hooks(self, workflow: StateGraph) -> str: """Add pre/post model hook nodes and return entrypoint.""" - # Implementation will be added in next task - pass + # Add pre_model_hook node with edge to agent if provided + if self.pre_model_hook is not None: + workflow.add_node("pre_model_hook", self.pre_model_hook) # type: ignore[arg-type] + workflow.add_edge("pre_model_hook", "agent") + entrypoint = "pre_model_hook" + else: + entrypoint = "agent" + + # Add post_model_hook node with conditional routing if provided + if self.post_model_hook is not None: + workflow.add_node("post_model_hook", self.post_model_hook) # type: ignore[arg-type] + + 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 + ] + + # Handle pending tool calls + 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 + ] + # Structured response generation + elif isinstance(messages[-1], ToolMessage): + return entrypoint + elif self.response_format is not None: + return "generate_structured_response" + # END routing based on state + else: + return END + + # Store the router function for use in build method + self.post_model_hook_router = post_model_hook_router + + return entrypoint def build(self) -> CompiledStateGraph: """Assemble the complete graph based on all configuration options.""" @@ -1418,3 +1473,4 @@ __all__ = [ +