From ceda61fe2ae124cbad3503a5301a15b009f4aea9 Mon Sep 17 00:00:00 2001 From: "open-swe[bot]" Date: Mon, 28 Jul 2025 20:50:20 +0000 Subject: [PATCH] Apply patch --- .../langgraph/prebuilt/chat_agent_executor.py | 139 +++++++++++++++++- 1 file changed, 137 insertions(+), 2 deletions(-) diff --git a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py index 01f6b24bd..22baaea53 100644 --- a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py +++ b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py @@ -427,8 +427,142 @@ class _AgentBuilder: def _create_model_node(self) -> RunnableCallable: """Create the core LLM interaction node.""" - # Implementation will be added in next task - pass + def _are_more_steps_needed(state: StateSchema, response: BaseMessage) -> bool: + has_tool_calls = isinstance(response, AIMessage) and response.tool_calls + all_tools_return_direct = ( + all(call["name"] in self.should_return_direct for call in response.tool_calls) + if isinstance(response, AIMessage) + else False + ) + remaining_steps = _get_state_value(state, "remaining_steps", None) + is_last_step = _get_state_value(state, "is_last_step", False) + return ( + (remaining_steps is None and is_last_step and has_tool_calls) + or ( + remaining_steps is not None + and remaining_steps < 1 + and all_tools_return_direct + ) + or (remaining_steps is not None and remaining_steps < 2 and has_tool_calls) + ) + + def _get_model_input_state(state: StateSchema) -> StateSchema: + if self.pre_model_hook is not None: + messages = ( + _get_state_value(state, "llm_input_messages") + ) or _get_state_value(state, "messages") + error_msg = f"Expected input to call_model to have 'llm_input_messages' or 'messages' key, but got {state}" + else: + messages = _get_state_value(state, "messages") + error_msg = ( + f"Expected input to call_model to have 'messages' key, but got {state}" + ) + + if messages is None: + raise ValueError(error_msg) + + # Message validation with _validate_chat_history + _validate_chat_history(messages) + + # we're passing messages under `messages` key, as this is expected by the prompt + if isinstance(self.state_schema, type) and issubclass(self.state_schema, BaseModel): + state.messages = messages # type: ignore + else: + state["messages"] = messages # type: ignore + + return state + + # Define the function that calls the model + def call_model( + state: StateSchema, runtime: Runtime[ContextT], config: RunnableConfig + ) -> StateSchema: + if self.is_async_dynamic_model: + msg = ( + "Async model callable provided but agent invoked synchronously. " + "Use agent.ainvoke() or agent.astream(), or " + "provide a sync model callable." + ) + raise RuntimeError(msg) + + model_input = _get_model_input_state(state) + + # Model resolution + if self.is_dynamic_model: + # Resolve dynamic model at runtime and apply prompt + dynamic_model = self._resolve_model(state, runtime) + response = cast(AIMessage, dynamic_model.invoke(model_input, config)) # type: ignore[arg-type] + else: + response = cast(AIMessage, self.static_model.invoke(model_input, config)) # type: ignore[union-attr] + + # add agent name to the AIMessage + response.name = self.name + + # Remaining steps management + if _are_more_steps_needed(state, response): + 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: StateSchema, runtime: Runtime[ContextT], config: RunnableConfig + ) -> StateSchema: + model_input = _get_model_input_state(state) + + # Model resolution + if self.is_dynamic_model: + # Resolve dynamic model at runtime and apply prompt + # (supports both sync and async) + dynamic_model = await self._aresolve_model(state, runtime) + response = cast(AIMessage, await dynamic_model.ainvoke(model_input, config)) # type: ignore[arg-type] + else: + response = cast(AIMessage, await self.static_model.ainvoke(model_input, config)) # type: ignore[union-attr] + + # add agent name to the AIMessage + response.name = self.name + + # Remaining steps management + if _are_more_steps_needed(state, response): + 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]} + + # Proper input schema handling + input_schema: StateSchemaType + if self.pre_model_hook is not None: + # Dynamically create a schema that inherits from state_schema and adds 'llm_input_messages' + if isinstance(self.state_schema, type) and issubclass(self.state_schema, BaseModel): + # For Pydantic schemas + from pydantic import create_model + + input_schema = create_model( + "CallModelInputSchema", + llm_input_messages=(list[AnyMessage], ...), + __base__=self.state_schema, + ) + else: + # For TypedDict schemas + class CallModelInputSchema(self.state_schema): # type: ignore + llm_input_messages: list[AnyMessage] + + input_schema = CallModelInputSchema + else: + input_schema = self.state_schema + + return RunnableCallable(call_model, acall_model, input_schema=input_schema) def _create_structured_response_node(self) -> Optional[RunnableCallable]: """Create structured output generation node if needed.""" @@ -1163,3 +1297,4 @@ __all__ = [ +