Apply patch

libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py

@@ -111,42 +111,6 @@ def _get_state_value(state: StateSchema, key: str, default: Any = None) -> Any:
     )
 
 
-def _get_prompt_runnable(prompt: Optional[Prompt]) -> 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: BaseMessage = SystemMessage(content=prompt)
-        prompt_runnable = RunnableCallable(
-            lambda state: [_system_message] + _get_state_value(state, "messages"),
-            name=PROMPT_RUNNABLE_NAME,
-        )
-    elif isinstance(prompt, SystemMessage):
-        prompt_runnable = RunnableCallable(
-            lambda state: [prompt] + _get_state_value(state, "messages"),
-            name=PROMPT_RUNNABLE_NAME,
-        )
-    elif inspect.iscoroutinefunction(prompt):
-        prompt_runnable = RunnableCallable(
-            None,
-            prompt,
-            name=PROMPT_RUNNABLE_NAME,
-        )
-    elif callable(prompt):
-        prompt_runnable = RunnableCallable(
-            prompt,
-            name=PROMPT_RUNNABLE_NAME,
-        )
-    elif isinstance(prompt, Runnable):
-        prompt_runnable = prompt
-    else:
-        raise ValueError(f"Got unexpected type for `prompt`: {type(prompt)}")
-
-    return prompt_runnable
-
-
 def _should_bind_tools(
     model: LanguageModelLike, tools: Sequence[BaseTool], num_builtin: int = 0
 ) -> bool:
@@ -248,6 +212,455 @@ def _validate_chat_history(
     raise ValueError(error_message)
 
 
+class _AgentBuilder:
+    """Internal builder class for constructing ReAct-style agent graphs."""
+
+    def __init__(
+        self,
+        model: Union[
+            str,
+            LanguageModelLike,
+            Callable[[StateSchema, Runtime[ContextT]], BaseChatModel],
+            Callable[[StateSchema, Runtime[ContextT]], Awaitable[BaseChatModel]],
+        ],
+        tools: Union[Sequence[Union[BaseTool, Callable, dict[str, Any]]], ToolNode],
+        *,
+        prompt: Optional[Prompt] = None,
+        response_format: Optional[
+            Union[StructuredResponseSchema, tuple[str, StructuredResponseSchema]]
+        ] = None,
+        pre_model_hook: Optional[RunnableLike] = None,
+        post_model_hook: Optional[RunnableLike] = None,
+        state_schema: Optional[StateSchemaType] = None,
+        context_schema: Optional[Type[Any]] = None,
+        checkpointer: Optional[Checkpointer] = None,
+        store: Optional[BaseStore] = None,
+        interrupt_before: Optional[list[str]] = None,
+        interrupt_after: Optional[list[str]] = None,
+        debug: bool = False,
+        version: Literal["v1", "v2"] = "v2",
+        name: Optional[str] = None,
+    ):
+        """Initialize the agent builder with all parameters."""
+        # Store all parameters
+        self.model = model
+        self.tools = tools
+        self.prompt = prompt
+        self.response_format = response_format
+        self.pre_model_hook = pre_model_hook
+        self.post_model_hook = post_model_hook
+        self.state_schema = state_schema
+        self.context_schema = context_schema
+        self.checkpointer = checkpointer
+        self.store = store
+        self.interrupt_before = interrupt_before
+        self.interrupt_after = interrupt_after
+        self.debug = debug
+        self.version = version
+        self.name = name
+
+        # Initialize derived attributes
+        self.tool_node: ToolNode
+        self.tool_classes: list[Type[BaseTool]]
+        self.llm_builtin_tools: list[dict] = []
+        self.is_dynamic_model: bool
+        self.is_async_dynamic_model: bool
+        self.tool_calling_enabled: bool
+        self.static_model: Optional[Runnable] = None
+        self.should_return_direct: set[str] = set()
+
+        # Run validation and setup
+        self._validate_version()
+        self._validate_state_schema()
+        self._validate_and_setup_tools()
+        self._validate_and_setup_model()
+
+    def _validate_version(self) -> None:
+        """Validate the version parameter."""
+        if self.version not in ("v1", "v2"):
+            raise ValueError(
+                f"Invalid version {self.version}. Supported versions are 'v1' and 'v2'."
+            )
+
+    def _validate_state_schema(self) -> None:
+        """Validate and set up the state schema."""
+        if self.state_schema is not None:
+            required_keys = {"messages", "remaining_steps"}
+            if self.response_format is not None:
+                required_keys.add("structured_response")
+
+            schema_keys = set(get_type_hints(self.state_schema))
+            if missing_keys := required_keys - set(schema_keys):
+                raise ValueError(f"Missing required key(s) {missing_keys} in state_schema")
+
+        if self.state_schema is None:
+            self.state_schema = (
+                AgentStateWithStructuredResponse
+                if self.response_format is not None
+                else AgentState
+            )
+
+    def _validate_and_setup_tools(self) -> None:
+        """Validate and set up tools and tool node."""
+        if isinstance(self.tools, ToolNode):
+            self.tool_classes = list(self.tools.tools_by_name.values())
+            self.tool_node = self.tools
+        else:
+            self.llm_builtin_tools = [t for t in self.tools if isinstance(t, dict)]
+            self.tool_node = ToolNode([t for t in self.tools if not isinstance(t, dict)])
+            self.tool_classes = list(self.tool_node.tools_by_name.values())
+
+        self.tool_calling_enabled = len(self.tool_classes) > 0
+
+        # If any of the tools are configured to return_directly after running,
+        # our graph needs to check if these were called
+        self.should_return_direct = {t.name for t in self.tool_classes if t.return_direct}
+
+    def _validate_and_setup_model(self) -> None:
+        """Validate and set up the model."""
+        self.is_dynamic_model = not isinstance(self.model, (str, Runnable)) and callable(self.model)
+        self.is_async_dynamic_model = self.is_dynamic_model and inspect.iscoroutinefunction(self.model)
+
+        if not self.is_dynamic_model:
+            model = self.model
+            if isinstance(model, str):
+                try:
+                    from langchain.chat_models import (  # type: ignore[import-not-found]
+                        init_chat_model,
+                    )
+                except ImportError:
+                    raise ImportError(
+                        "Please install langchain (`pip install langchain`) to "
+                        "use '<provider>:<model>' string syntax for `model` parameter."
+                    )
+
+                model = cast(BaseChatModel, init_chat_model(model))
+
+            if (
+                _should_bind_tools(model, self.tool_classes, num_builtin=len(self.llm_builtin_tools))  # type: ignore[arg-type]
+                and len(self.tool_classes + self.llm_builtin_tools) > 0
+            ):
+                model = cast(BaseChatModel, model).bind_tools(
+                    self.tool_classes + self.llm_builtin_tools  # type: ignore[operator]
+                )
+
+            self.static_model = self._get_prompt_runnable() | model  # type: ignore[operator]
+        else:
+            # For dynamic models, we'll create the runnable at runtime
+            self.static_model = None
+
+    def _get_prompt_runnable(self) -> Runnable:
+        """Get the prompt runnable based on the prompt configuration."""
+        prompt_runnable: Runnable
+        if self.prompt is None:
+            prompt_runnable = RunnableCallable(
+                lambda state: _get_state_value(state, "messages"), name=PROMPT_RUNNABLE_NAME
+            )
+        elif isinstance(self.prompt, str):
+            _system_message: BaseMessage = SystemMessage(content=self.prompt)
+            prompt_runnable = RunnableCallable(
+                lambda state: [_system_message] + _get_state_value(state, "messages"),
+                name=PROMPT_RUNNABLE_NAME,
+            )
+        elif isinstance(self.prompt, SystemMessage):
+            prompt_runnable = RunnableCallable(
+                lambda state: [self.prompt] + _get_state_value(state, "messages"),
+                name=PROMPT_RUNNABLE_NAME,
+            )
+        elif inspect.iscoroutinefunction(self.prompt):
+            prompt_runnable = RunnableCallable(
+                None,
+                self.prompt,
+                name=PROMPT_RUNNABLE_NAME,
+            )
+        elif callable(self.prompt):
+            prompt_runnable = RunnableCallable(
+                self.prompt,
+                name=PROMPT_RUNNABLE_NAME,
+            )
+        elif isinstance(self.prompt, Runnable):
+            prompt_runnable = self.prompt
+        else:
+            raise ValueError(f"Got unexpected type for `prompt`: {type(self.prompt)}")
+
+        return prompt_runnable
+
+    def _resolve_model(
+        self, state: StateSchema, runtime: Runtime[ContextT]
+    ) -> LanguageModelLike:
+        """Resolve the model to use, handling both static and dynamic models."""
+        if self.is_dynamic_model:
+            return self._get_prompt_runnable() | self.model(state, runtime)  # type: ignore[operator]
+        else:
+            return self.static_model
+
+    async def _aresolve_model(
+        self, state: StateSchema, runtime: Runtime[ContextT]
+    ) -> LanguageModelLike:
+        """Async resolve the model to use, handling both static and dynamic models."""
+        if self.is_async_dynamic_model:
+            resolved_model = await self.model(state, runtime)  # type: ignore[misc,operator]
+            return self._get_prompt_runnable() | resolved_model
+        elif self.is_dynamic_model:
+            return self._get_prompt_runnable() | self.model(state, runtime)  # type: ignore[operator]
+        else:
+            return self.static_model
+
+    def _are_more_steps_needed(self, state: StateSchema, response: BaseMessage) -> bool:
+        """Check if more steps are needed based on the response and state."""
+        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(self, state: StateSchema) -> StateSchema:
+        """Get the model input state, handling pre_model_hook if present."""
+        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)
+
+        _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
+
+    def _create_model_node(self) -> tuple[Callable, Callable]:
+        """Create the model node functions (call_model and acall_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 = self._get_model_input_state(state)
+
+            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
+
+            if self._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 = self._get_model_input_state(state)
+
+            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
+            if self._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]}
+
+        return call_model, acall_model
+
+    def _create_structured_response_node(self) -> tuple[Callable, Callable]:
+        """Create the structured response node functions."""
+        def generate_structured_response(
+            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)
+
+            messages = _get_state_value(state, "messages")
+            structured_response_schema = self.response_format
+            if isinstance(self.response_format, tuple):
+                system_prompt, structured_response_schema = self.response_format
+                messages = [SystemMessage(content=system_prompt)] + list(messages)
+
+            resolved_model = self._resolve_model(state, runtime)
+            model_with_structured_output = _get_model(
+                resolved_model
+            ).with_structured_output(
+                cast(StructuredResponseSchema, structured_response_schema)
+            )
+            response = model_with_structured_output.invoke(messages, config)
+            return {"structured_response": response}
+
+        async def agenerate_structured_response(
+            state: StateSchema, runtime: Runtime[ContextT], config: RunnableConfig
+        ) -> StateSchema:
+            messages = _get_state_value(state, "messages")
+            structured_response_schema = self.response_format
+            if isinstance(self.response_format, tuple):
+                system_prompt, structured_response_schema = self.response_format
+                messages = [SystemMessage(content=system_prompt)] + list(messages)
+
+            resolved_model = await self._aresolve_model(state, runtime)
+            model_with_structured_output = _get_model(
+                resolved_model
+            ).with_structured_output(
+                cast(StructuredResponseSchema, structured_response_schema)
+            )
+            response = await model_with_structured_output.ainvoke(messages, config)
+            return {"structured_response": response}
+
+        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[
         str,
@@ -952,3 +1365,8 @@ __all__ = [
     "AgentStateWithStructuredResponse",
     "AgentStateWithStructuredResponsePydantic",
 ]
+
+
+
+
+