From e8dd6823205e90ffe12aa1a8437526a83b34ef0d Mon Sep 17 00:00:00 2001
From: Vadym Barda <vadym@langchain.dev>
Date: Wed, 5 Mar 2025 21:37:55 -0500
Subject: [PATCH] prebuilt: allow passing RunnableSequence as a model (#3706)

---
 .../langgraph/prebuilt/chat_agent_executor.py | 21 ++++++
 libs/prebuilt/tests/test_react_agent.py       | 65 +++++++++++++++++++
 2 files changed, 86 insertions(+)

diff --git a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py
index 5c34f1ecc..117a2bc91 100644
--- a/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py
+++ b/libs/prebuilt/langgraph/prebuilt/chat_agent_executor.py
@@ -22,6 +22,7 @@ from langchain_core.runnables import (
     Runnable,
     RunnableBinding,
     RunnableConfig,
+    RunnableSequence,
 )
 from langchain_core.tools import BaseTool
 from pydantic import BaseModel
@@ -133,6 +134,16 @@ def _convert_modifier_to_prompt(func: F) -> F:
 
 
 def _should_bind_tools(model: LanguageModelLike, tools: Sequence[BaseTool]) -> bool:
+    if isinstance(model, RunnableSequence):
+        model = next(
+            (
+                step
+                for step in model.steps
+                if isinstance(step, (RunnableBinding, BaseChatModel))
+            ),
+            model,
+        )
+
     if not isinstance(model, RunnableBinding):
         return True
 
@@ -168,6 +179,16 @@ def _should_bind_tools(model: LanguageModelLike, tools: Sequence[BaseTool]) -> b
 
 def _get_model(model: LanguageModelLike) -> BaseChatModel:
     """Get the underlying model from a RunnableBinding or return the model itself."""
+    if isinstance(model, RunnableSequence):
+        model = next(
+            (
+                step
+                for step in model.steps
+                if isinstance(step, (RunnableBinding, BaseChatModel))
+            ),
+            model,
+        )
+
     if isinstance(model, RunnableBinding):
         model = model.bound
 
diff --git a/libs/prebuilt/tests/test_react_agent.py b/libs/prebuilt/tests/test_react_agent.py
index 0b6f10f51..d5f405384 100644
--- a/libs/prebuilt/tests/test_react_agent.py
+++ b/libs/prebuilt/tests/test_react_agent.py
@@ -35,6 +35,8 @@ from langgraph.prebuilt import (
 )
 from langgraph.prebuilt.chat_agent_executor import (
     AgentState,
+    _get_model,
+    _should_bind_tools,
     _validate_chat_history,
 )
 from langgraph.prebuilt.tool_node import (
@@ -1324,3 +1326,66 @@ def test_tool_node_node_interrupt(
             ns=[AnyStr("tools:")],
         ),
     )
+
+
+@pytest.mark.parametrize("tool_style", ["openai", "anthropic"])
+def test_should_bind_tools(tool_style: str) -> None:
+    @dec_tool
+    def some_tool(some_val: int) -> str:
+        """Tool docstring."""
+        return "meow"
+
+    @dec_tool
+    def some_other_tool(some_val: int) -> str:
+        """Tool docstring."""
+        return "meow"
+
+    model = FakeToolCallingModel(tool_style=tool_style)
+    # should bind when a regular model
+    assert _should_bind_tools(model, [])
+    assert _should_bind_tools(model, [some_tool])
+
+    # should bind when a seq
+    seq = model | RunnableLambda(lambda message: message)
+    assert _should_bind_tools(seq, [])
+    assert _should_bind_tools(seq, [some_tool])
+
+    # should not bind when a model with tools
+    assert not _should_bind_tools(model.bind_tools([some_tool]), [some_tool])
+    # should not bind when a seq with tools
+    seq_with_tools = model.bind_tools([some_tool]) | RunnableLambda(
+        lambda message: message
+    )
+    assert not _should_bind_tools(seq_with_tools, [some_tool])
+
+    # should raise on invalid inputs
+    with pytest.raises(ValueError):
+        _should_bind_tools(model.bind_tools([some_tool]), [])
+    with pytest.raises(ValueError):
+        _should_bind_tools(model.bind_tools([some_tool]), [some_other_tool])
+    with pytest.raises(ValueError):
+        _should_bind_tools(model.bind_tools([some_tool]), [some_tool, some_other_tool])
+
+
+def test_get_model() -> None:
+    model = FakeToolCallingModel(tool_calls=[])
+    assert _get_model(model) == model
+
+    @dec_tool
+    def some_tool(some_val: int) -> str:
+        """Tool docstring."""
+        return "meow"
+
+    model_with_tools = model.bind_tools([some_tool])
+    assert _get_model(model_with_tools) == model
+
+    seq = model | RunnableLambda(lambda message: message)
+    assert _get_model(seq) == model
+
+    seq_with_tools = model.bind_tools([some_tool]) | RunnableLambda(
+        lambda message: message
+    )
+    assert _get_model(seq_with_tools) == model
+
+    with pytest.raises(TypeError):
+        _get_model(RunnableLambda(lambda message: message))