From d5a835e5fd4eb20b9a5ee86e84622cced9aebbe7 Mon Sep 17 00:00:00 2001 From: Eugene Yurtsev Date: Wed, 13 Aug 2025 10:27:56 -0400 Subject: [PATCH] x --- .../langgraph/prebuilt/chat_agent_executor.py | 71 +++++++++++++++++-- 1 file changed, 67 insertions(+), 4 deletions(-) diff --git a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py index f7df559ee..32faaf317 100644 --- a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py +++ b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py @@ -55,6 +55,12 @@ StructuredResponseSchema = Union[dict, type[BaseModel]] F = TypeVar("F", bound=Callable[..., Any]) +class StepCountIs(BaseModel): + """Stop condition that specifies when to halt agent execution based on step count.""" + + count: int + + # We create the AgentState that we will pass around # This simply involves a list of messages # We want steps to return messages to append to the list @@ -66,6 +72,8 @@ class AgentState(TypedDict): remaining_steps: NotRequired[RemainingSteps] + model_calls: NotRequired[int] + class AgentStatePydantic(BaseModel): """The state of the agent.""" @@ -74,6 +82,8 @@ class AgentStatePydantic(BaseModel): remaining_steps: RemainingSteps = 25 + model_calls: int = 0 + class AgentStateWithStructuredResponse(AgentState): """The state of the agent with a structured response.""" @@ -270,6 +280,7 @@ class _AgentBuilder: response_format: Optional[ Union[StructuredResponseSchema, tuple[str, StructuredResponseSchema]] ] = None, + stop_when: Optional[Union[StepCountIs, Callable[[StateSchema], bool]]] = None, pre_model_hook: Optional[RunnableLike] = None, post_model_hook: Optional[RunnableLike] = None, state_schema: Optional[StateSchemaType] = None, @@ -291,6 +302,7 @@ class _AgentBuilder: self.tools = tools self.prompt = prompt self.response_format = response_format + self.stop_when = stop_when self.pre_model_hook = pre_model_hook self.post_model_hook = post_model_hook self.state_schema = state_schema @@ -328,6 +340,8 @@ class _AgentBuilder: required_keys = {"messages", "remaining_steps"} if self.response_format is not None: required_keys.add("structured_response") + if self.stop_when is not None: + required_keys.add("model_calls") schema_keys = set(get_type_hints(self.state_schema)) if missing_keys := required_keys - schema_keys: @@ -459,6 +473,19 @@ class _AgentBuilder: return True return False + def _should_stop_execution(state: StateSchema) -> bool: + """Check if execution should stop based on stop_when condition.""" + if self.stop_when is None: + return False + + if isinstance(self.stop_when, StepCountIs): + model_calls = _get_state_value(state, "model_calls", 0) + return model_calls >= self.stop_when.count + elif callable(self.stop_when): + return self.stop_when(state) + + return False + def call_model( state: StateSchema, runtime: Runtime[ContextT], config: RunnableConfig ) -> StateSchema: @@ -468,25 +495,40 @@ class _AgentBuilder: "Use agent.ainvoke() or agent.astream(), or provide a sync model callable." ) + # Check if we should stop before making the model call + if _should_stop_execution(state): + return {} + model_input = _get_model_input_state(state) model = self._resolve_model(state, runtime) response = cast(AIMessage, model.invoke(model_input, config)) # type: ignore[arg-type] response.name = self.name + # Track model calls + result = {"messages": [response]} + if self.stop_when is not None: + current_calls = _get_state_value(state, "model_calls", 0) + result["model_calls"] = current_calls + 1 + if _are_more_steps_needed(state, response): return { + **result, "messages": [ AIMessage( id=response.id, content="Sorry, need more steps to process this request.", ) - ] + ], } - return {"messages": [response]} + return result async def acall_model( state: StateSchema, runtime: Runtime[ContextT], config: RunnableConfig ) -> StateSchema: + # Check if we should stop before making the model call + if _should_stop_execution(state): + return {} + model_input = _get_model_input_state(state) model = await self._aresolve_model(state, runtime) @@ -495,16 +537,24 @@ class _AgentBuilder: await model.ainvoke(model_input, config), # type: ignore[arg-type] ) response.name = self.name + + # Track model calls + result = {"messages": [response]} + if self.stop_when is not None: + current_calls = _get_state_value(state, "model_calls", 0) + result["model_calls"] = current_calls + 1 + if _are_more_steps_needed(state, response): return { + **result, "messages": [ AIMessage( id=response.id, content="Sorry, need more steps to process this request.", ) - ] + ], } - return {"messages": [response]} + return result return RunnableCallable(call_model, acall_model) @@ -852,6 +902,7 @@ def create_react_agent( response_format: Optional[ Union[StructuredResponseSchema, tuple[str, StructuredResponseSchema]] ] = None, + stop_when: Optional[Union[StepCountIs, Callable[[StateSchema], bool]]] = None, pre_model_hook: Optional[RunnableLike] = None, post_model_hook: Optional[RunnableLike] = None, state_schema: Optional[StateSchemaType] = None, @@ -940,6 +991,16 @@ def create_react_agent( The graph will make a separate call to the LLM to generate the structured response after the agent loop is finished. This is not the only strategy to get structured responses, see more options in [this guide](https://langchain-ai.github.io/langgraph/how-tos/react-agent-structured-output/). + stop_when: An optional condition to stop agent execution. + + Can be passed in as: + + - `StepCountIs(count=N)`: Stop execution after exactly N model calls + - A callable with signature `(state) -> bool`: Custom stop condition that returns True when execution should stop + + When a stop condition is met, the agent will halt further execution and return the current state. + If using `StepCountIs`, the state will include a `model_calls` field tracking the number of model invocations. + pre_model_hook: An optional node to add before the `agent` node (i.e., the node that calls the LLM). Useful for managing long message histories (e.g., message trimming, summarization, etc.). Pre-model hook must be a callable or a runnable that takes in current graph state and returns a state update in the form of @@ -1081,6 +1142,7 @@ def create_react_agent( tools=tools, prompt=prompt, response_format=response_format, + stop_when=stop_when, pre_model_hook=pre_model_hook, post_model_hook=post_model_hook, state_schema=state_schema, @@ -1113,4 +1175,5 @@ __all__ = [ "AgentStatePydantic", "AgentStateWithStructuredResponse", "AgentStateWithStructuredResponsePydantic", + "StepCountIs", ]