chore: migrate legacy langgraph imports to maintain compatibility

PiperOrigin-RevId: 906317678

Author: vertex-sdk-bot

tests/unit/vertex_langchain/test_agent_engine_templates_langgraph.py

@@ -208,7 +208,7 @@ def test_query(self, langchain_dump_mock):
         mocks.attach_mock(mock=agent._tmpl_attrs.get("runnable"), attribute="invoke")
         agent.query(input="test query")
         mocks.assert_has_calls(
-            [mock.call.invoke.invoke(input={"input": "test query"}, config=None)]
+            [mock.call.invoke.invoke(input={"input": "test query", "messages": [("user", "test query")]}, config=None)]
         )
 
     def test_stream_query(self, langchain_dump_mock):
@@ -217,7 +217,7 @@ def test_stream_query(self, langchain_dump_mock):
         agent._tmpl_attrs["runnable"].stream.return_value = []
         list(agent.stream_query(input="test stream query"))
         agent._tmpl_attrs["runnable"].stream.assert_called_once_with(
-            input={"input": "test stream query"},
+            input={"input": "test stream query", "messages": [("user", "test stream query")]},
             config=None,
         )
 

tests/unit/vertex_langchain/test_reasoning_engine_templates_langgraph.py

@@ -208,7 +208,7 @@ def test_query(self, langchain_dump_mock):
         mocks.attach_mock(mock=agent._runnable, attribute="invoke")
         agent.query(input="test query")
         mocks.assert_has_calls(
-            [mock.call.invoke.invoke(input={"input": "test query"}, config=None)]
+            [mock.call.invoke.invoke(input={"input": "test query", "messages": [("user", "test query")]}, config=None)]
         )
 
     def test_stream_query(self, langchain_dump_mock):
@@ -217,7 +217,7 @@ def test_stream_query(self, langchain_dump_mock):
         agent._runnable.stream.return_value = []
         list(agent.stream_query(input="test stream query"))
         agent._runnable.stream.assert_called_once_with(
-            input={"input": "test stream query"},
+            input={"input": "test stream query", "messages": [("user", "test stream query")]},
             config=None,
         )
 

vertexai/agent_engines/templates/langchain.py

@@ -43,11 +43,12 @@
         RunnableSerializable = Any
 
     try:
-        from langchain_google_vertexai.functions_utils import _ToolsType
-
-        _ToolsType = _ToolsType
+        from langchain_google_genai.functions_utils import _ToolsType
     except ImportError:
-        _ToolsType = Any
+        try:
+            from langchain_google_vertexai.functions_utils import _ToolsType
+        except ImportError:
+            _ToolsType = Any
 
     try:
         from opentelemetry.sdk import trace
@@ -81,13 +82,15 @@ def _default_runnable_kwargs(has_history: bool) -> Mapping[str, Any]:
 
 def _default_output_parser():
     try:
-        from langchain.agents.output_parsers.tools import ToolsAgentOutputParser
+        from langchain_classic.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,
-        )
-
+        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()
 
 
@@ -98,17 +101,29 @@ def _default_model_builder(
     location: str,
     model_kwargs: Optional[Mapping[str, Any]] = None,
 ) -> "BaseLanguageModel":
-    import vertexai
-    from google.cloud.aiplatform import initializer
-    from langchain_google_vertexai import ChatVertexAI
-
     model_kwargs = model_kwargs or {}
-    current_project = initializer.global_config.project
-    current_location = initializer.global_config.location
-    vertexai.init(project=project, location=location)
-    model = ChatVertexAI(model_name=model_name, **model_kwargs)
-    vertexai.init(project=current_project, location=current_location)
-    return model
+    try:
+        from langchain_google_genai import ChatGoogleGenerativeAI
+
+        model = ChatGoogleGenerativeAI(
+            model=model_name,
+            project=project,
+            location=location,
+            vertexai=True,
+            **model_kwargs,
+        )
+        return model
+    except ImportError:
+        import vertexai
+        from google.cloud.aiplatform import initializer
+        from langchain_google_vertexai import ChatVertexAI
+
+        current_project = initializer.global_config.project
+        current_location = initializer.global_config.location
+        vertexai.init(project=project, location=location)
+        model = ChatVertexAI(model_name=model_name, **model_kwargs)
+        vertexai.init(project=current_project, location=current_location)
+        return model
 
 
 def _default_runnable_builder(
@@ -124,8 +139,16 @@ def _default_runnable_builder(
     runnable_kwargs: Optional[Mapping[str, Any]] = None,
 ) -> "RunnableSerializable":
     from langchain_core import tools as lc_tools
-    from langchain.agents import AgentExecutor
-    from langchain.tools.base import StructuredTool
+
+    try:
+        from langchain_classic.agents import AgentExecutor
+    except ImportError:
+        from langchain.agents import AgentExecutor
+
+    try:
+        from langchain_core.tools import StructuredTool
+    except ImportError:
+        from langchain.tools.base import StructuredTool
 
     # The prompt template and runnable_kwargs needs to be customized depending
     # on whether the user intends for the agent to have history. The way the
@@ -261,12 +284,16 @@ def _default_prompt(
     from langchain_core import prompts
 
     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,
+        from langchain_classic.agents.format_scratchpad.tools import (
+            format_to_tool_messages,
         )
+    except (ModuleNotFoundError, ImportError):
+        try:
+            from langchain.agents.format_scratchpad.tools import format_to_tool_messages
+        except (ModuleNotFoundError, ImportError):
+            from langchain.agents.format_scratchpad.openai_tools import (
+                format_to_openai_tool_messages as format_to_tool_messages,
+            )
 
     system_instructions = []
     if system_instruction:
@@ -629,13 +656,18 @@ def query(
         Returns:
             The output of querying the Agent with the given input and config.
         """
-        from langchain.load import dump as langchain_load_dump
+        try:
+            from langchain_core.load import dumpd
+        except ImportError:
+            from langchain.load import dump as langchain_load_dump
+
+            dumpd = langchain_load_dump.dumpd
 
         if isinstance(input, str):
             input = {"input": input}
         if not self._tmpl_attrs.get("runnable"):
             self.set_up()
-        return langchain_load_dump.dumpd(
+        return dumpd(
             self._tmpl_attrs.get("runnable").invoke(
                 input=input, config=config, **kwargs
             )
@@ -662,7 +694,12 @@ def stream_query(
         Yields:
             The output of querying the Agent with the given input and config.
         """
-        from langchain.load import dump as langchain_load_dump
+        try:
+            from langchain_core.load import dumpd
+        except ImportError:
+            from langchain.load import dump as langchain_load_dump
+
+            dumpd = langchain_load_dump.dumpd
 
         if isinstance(input, str):
             input = {"input": input}
@@ -673,4 +710,4 @@ def stream_query(
             config=config,
             **kwargs,
         ):
-            yield langchain_load_dump.dumpd(chunk)
+            yield dumpd(chunk)

vertexai/agent_engines/templates/langgraph.py

@@ -34,11 +34,14 @@
         BaseLanguageModel = Any
 
     try:
-        from langchain_google_vertexai.functions_utils import _ToolsType
-
+        from langchain_google_genai.functions_utils import _ToolsType
         _ToolLike = _ToolsType
     except ImportError:
-        _ToolLike = Any
+        try:
+            from langchain_google_vertexai.functions_utils import _ToolsType
+            _ToolLike = _ToolsType
+        except ImportError:
+            _ToolLike = Any
 
     try:
         from opentelemetry.sdk import trace
@@ -53,12 +56,10 @@
 
     try:
         from langgraph_checkpoint.checkpoint import base
-
         BaseCheckpointSaver = base.BaseCheckpointSaver
     except ImportError:
         try:
             from langgraph.checkpoint import base
-
             BaseCheckpointSaver = base.BaseCheckpointSaver
         except ImportError:
             BaseCheckpointSaver = Any
@@ -87,17 +88,29 @@ def _default_model_builder(
     Returns:
         BaseLanguageModel: The language model.
     """
-    import vertexai
-    from google.cloud.aiplatform import initializer
-    from langchain_google_vertexai import ChatVertexAI
-
     model_kwargs = model_kwargs or {}
-    current_project = initializer.global_config.project
-    current_location = initializer.global_config.location
-    vertexai.init(project=project, location=location)
-    model = ChatVertexAI(model_name=model_name, **model_kwargs)
-    vertexai.init(project=current_project, location=current_location)
-    return model
+    try:
+        from langchain_google_genai import ChatGoogleGenerativeAI
+
+        model = ChatGoogleGenerativeAI(
+            model=model_name,
+            project=project,
+            location=location,
+            vertexai=True,
+            **model_kwargs,
+        )
+        return model
+    except ImportError:
+        import vertexai
+        from google.cloud.aiplatform import initializer
+        from langchain_google_vertexai import ChatVertexAI
+
+        current_project = initializer.global_config.project
+        current_location = initializer.global_config.location
+        vertexai.init(project=project, location=location)
+        model = ChatVertexAI(model_name=model_name, **model_kwargs)
+        vertexai.init(project=current_project, location=current_location)
+        return model
 
 
 def _default_runnable_builder(
@@ -554,13 +567,16 @@ def query(
         Returns:
             The output of querying the Agent with the given input and config.
         """
-        from langchain.load import dump as langchain_load_dump
+        try:
+            from langchain_core.load import dumpd
+        except ImportError:
+            from langchain.load.dump import dumpd
 
         if isinstance(input, str):
-            input = {"input": input}
+            input = {"input": input, "messages": [("user", input)]}
         if not self._tmpl_attrs.get("runnable"):
             self.set_up()
-        return langchain_load_dump.dumpd(
+        return dumpd(
             self._tmpl_attrs.get("runnable").invoke(
                 input=input, config=config, **kwargs
             )
@@ -587,18 +603,21 @@ def stream_query(
         Yields:
             The output of querying the Agent with the given input and config.
         """
-        from langchain.load import dump as langchain_load_dump
+        try:
+            from langchain_core.load import dumpd
+        except ImportError:
+            from langchain.load.dump import dumpd
 
         if isinstance(input, str):
-            input = {"input": input}
+            input = {"input": input, "messages": [("user", input)]}
         if not self._tmpl_attrs.get("runnable"):
             self.set_up()
         for chunk in self._tmpl_attrs.get("runnable").stream(
             input=input,
             config=config,
             **kwargs,
         ):
-            yield langchain_load_dump.dumpd(chunk)
+            yield dumpd(chunk)
 
     def get_state_history(
         self,