feat: Support VertexTool in langchain template.

PiperOrigin-RevId: 639027200

Author: yeesian

tests/unit/vertex_langchain/test_reasoning_engine_templates_langchain.py

@@ -21,6 +21,8 @@
 import vertexai
 from google.cloud.aiplatform import initializer
 from vertexai.preview import reasoning_engines
+from vertexai.preview.generative_models import grounding
+from vertexai.generative_models import Tool
 import pytest
 
 
@@ -81,6 +83,12 @@ def langchain_dump_mock():
         yield langchain_dump_mock
 
 
+@pytest.fixture
+def mock_chatvertexai():
+    with mock.patch("langchain_google_vertexai.ChatVertexAI") as model_mock:
+        yield model_mock
+
+
 @pytest.mark.usefixtures("google_auth_mock")
 class TestLangchainAgent:
     def setup_method(self):
@@ -113,19 +121,23 @@ def test_initialization(self):
         assert agent._location == _TEST_LOCATION
         assert agent._runnable is None
 
-    def test_initialization_with_tools(self):
+    def test_initialization_with_tools(self, mock_chatvertexai):
         tools = [
             place_tool_query,
             StructuredTool.from_function(place_photo_query),
+            Tool.from_google_search_retrieval(grounding.GoogleSearchRetrieval()),
         ]
         agent = reasoning_engines.LangchainAgent(
             model=_TEST_MODEL,
             tools=tools,
         )
         for tool, agent_tool in zip(tools, agent._tools):
             assert isinstance(agent_tool, type(tool))
+        assert agent._runnable is None
+        agent.set_up()
+        assert agent._runnable is not None
 
-    def test_set_up(self, vertexai_init_mock):
+    def test_set_up(self):
         agent = reasoning_engines.LangchainAgent(
             model=_TEST_MODEL,
             prompt=self.prompt,
@@ -135,7 +147,7 @@ def test_set_up(self, vertexai_init_mock):
         agent.set_up()
         assert agent._runnable is not None
 
-    def test_clone(self, vertexai_init_mock):
+    def test_clone(self):
         agent = reasoning_engines.LangchainAgent(
             model=_TEST_MODEL,
             prompt=self.prompt,

vertexai/preview/reasoning_engines/templates/langchain.py

@@ -42,6 +42,13 @@
         RunnableConfig = Any
         RunnableSerializable = Any
 
+    try:
+        from langchain_google_vertexai.functions_utils import _ToolsType
+
+        _ToolLike = _ToolsType
+    except ImportError:
+        _ToolLike = Any
+
 
 def _default_runnable_kwargs(has_history: bool) -> Mapping[str, Any]:
     # https://github.com/langchain-ai/langchain/blob/5784dfed001730530637793bea1795d9d5a7c244/libs/core/langchain_core/runnables/history.py#L237-L241
@@ -62,7 +69,13 @@ def _default_runnable_kwargs(has_history: bool) -> Mapping[str, Any]:
 
 
 def _default_output_parser():
-    from langchain.agents.output_parsers.tools import ToolsAgentOutputParser
+    try:
+        from langchain.agents.output_parsers.tools import ToolsAgentOutputParser
+    except (ModuleNotFoundError, ImportError):
+        # Fallback to an older version if needed.
+        from langchain.agents.output_parsers.openai_tools import (
+            OpenAIToolsAgentOutputParser as ToolsAgentOutputParser,
+        )
 
     return ToolsAgentOutputParser()
 
@@ -90,7 +103,7 @@ def _default_model_builder(
 def _default_runnable_builder(
     model: "BaseLanguageModel",
     *,
-    tools: Optional[Sequence[Union[Callable, "BaseTool"]]] = None,
+    tools: Optional[Sequence["_ToolLike"]] = None,
     prompt: Optional["RunnableSerializable"] = None,
     output_parser: Optional["RunnableSerializable"] = None,
     chat_history: Optional["GetSessionHistoryCallable"] = None,
@@ -123,6 +136,7 @@ def _default_runnable_builder(
             if isinstance(tool, lc_tools.BaseTool)
             else StructuredTool.from_function(tool)
             for tool in tools
+            if isinstance(tool, (Callable, lc_tools.BaseTool))
         ],
         **agent_executor_kwargs,
     )
@@ -139,7 +153,14 @@ def _default_runnable_builder(
 
 def _default_prompt(has_history: bool) -> "RunnableSerializable":
     from langchain_core import prompts
-    from langchain.agents.format_scratchpad.tools import format_to_tool_messages
+
+    try:
+        from langchain.agents.format_scratchpad.tools import format_to_tool_messages
+    except (ModuleNotFoundError, ImportError):
+        # Fallback to an older version if needed.
+        from langchain.agents.format_scratchpad.openai_tools import (
+            format_to_openai_tool_messages as format_to_tool_messages,
+        )
 
     if has_history:
         return {
@@ -186,12 +207,10 @@ def _validate_callable_parameters_are_annotated(callable: Callable):
             )
 
 
-def _validate_tools(tools: Sequence[Union[Callable, "BaseTool"]]):
+def _validate_tools(tools: Sequence["_ToolLike"]):
     """Validates that the tools are usable for tool calling."""
-    from langchain_core import tools as lc_tools
-
     for tool in tools:
-        if not isinstance(tool, lc_tools.BaseTool):
+        if isinstance(tool, Callable):
             _validate_callable_parameters_are_annotated(tool)
 
 
@@ -208,7 +227,7 @@ def __init__(
         model: str,
         *,
         prompt: Optional["RunnableSerializable"] = None,
-        tools: Optional[Sequence[Union[Callable, "BaseTool"]]] = None,
+        tools: Optional[Sequence["_ToolLike"]] = None,
         output_parser: Optional["RunnableSerializable"] = None,
         chat_history: Optional["GetSessionHistoryCallable"] = None,
         model_kwargs: Optional[Mapping[str, Any]] = None,