feat: Add support for BaseModels in LangChain templates

yeesian · copybara-github · commit 5eb885ee7e01 · 2024-05-02T18:04:42.000-07:00
PiperOrigin-RevId: 630232213
diff --git a/.release-please-manifest.json b/.release-please-manifest.json
@@ -1,3 +1,3 @@
 {
-    ".": "1.50.0"
+    ".": "1.49.0"
 }
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,29 +1,5 @@
 # Changelog
 
-## [1.50.0](https://github.com/googleapis/python-aiplatform/compare/v1.49.0...v1.50.0) (2024-05-02)
-
-
-### Features
-
-* Add `Candidate.grounding_metadata` property ([b22a8b8](https://github.com/googleapis/python-aiplatform/commit/b22a8b847e3b299b828e37405e3678093486de28))
-* Add option to not include time_series_metrics in get_experiment_df call. This will improve execution time for Experiments with large number of runs. ([78a95c5](https://github.com/googleapis/python-aiplatform/commit/78a95c52d0e7bd9ec5b656ce67044b2f01677156))
-* Add tune_model and deploy_tuned_model for TextEmbeddingModel. ([42f5d6f](https://github.com/googleapis/python-aiplatform/commit/42f5d6f7cd13d51c4a73113c59e8b3c728cfc08b))
-* Automatically populate parents for full resource name in Vertex RAG SDK ([26657ff](https://github.com/googleapis/python-aiplatform/commit/26657ffd25ecb91882ca764e513c2e952833257f))
-* Deploy a tuned text embedding model -- it doesn't matter, if it's tuned using Node.js, or curl. ([8ca9cdf](https://github.com/googleapis/python-aiplatform/commit/8ca9cdf3576e3ce3b373ace4cd6ab0e9c54aa9f2))
-* Make get_embeddings work both for foundational & tuned models. ([b8b589c](https://github.com/googleapis/python-aiplatform/commit/b8b589ce9fff29d1721450d32b4a84a7f69413c3))
-* Python SDK for Vertex Model Monitoring V2. ([021d59f](https://github.com/googleapis/python-aiplatform/commit/021d59f1487e4e16c847d4135899d6845c0210aa))
-* Support public endpoint for Ray Client ([57a5f78](https://github.com/googleapis/python-aiplatform/commit/57a5f7815ffb8523e91d900da4ff7cfd0c344fe4))
-
-
-### Bug Fixes
-
-* Add deprecation warnings when using Ray v2.4 ([3a36784](https://github.com/googleapis/python-aiplatform/commit/3a367843840513e3257610c8ab38e9f79d3bcea0))
-* Append allowed_plugins in tb-gcp-uploader to default allowed plugins ([aab9c3e](https://github.com/googleapis/python-aiplatform/commit/aab9c3e41b92a1d60090e3d1d594390a5e9f3ff6))
-* LLM - Added missing parameters to the no-op `_TunableTextEmbeddingModelMixin.get_tuned_model` method ([eb05ac4](https://github.com/googleapis/python-aiplatform/commit/eb05ac421f186441a92c6e3b6a010d74caf14782))
-* LVM - Fixed the typo in the VisionModel aspect ratio type annotation ([2d19137](https://github.com/googleapis/python-aiplatform/commit/2d1913773cf9f4a4f8a2c8c8f45680c3ea97f68e))
-* Move torch import ([e6d34df](https://github.com/googleapis/python-aiplatform/commit/e6d34df7da7508c655eb17ee694e1ab2160fc8aa))
-* Ray - Fixed exception when using Ray 2.4 ([2661f52](https://github.com/googleapis/python-aiplatform/commit/2661f52fd08169e5d29b58f2afce9702b30101ae))
-
 ## [1.49.0](https://github.com/googleapis/python-aiplatform/compare/v1.48.0...v1.49.0) (2024-04-27)
 
 
diff --git a/google/cloud/aiplatform/gapic_version.py b/google/cloud/aiplatform/gapic_version.py
@@ -13,4 +13,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-__version__ = "1.50.0"  # {x-release-please-version}
+__version__ = "1.49.0"  # {x-release-please-version}
diff --git a/google/cloud/aiplatform/v1/schema/predict/instance/gapic_version.py b/google/cloud/aiplatform/v1/schema/predict/instance/gapic_version.py
@@ -13,4 +13,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-__version__ = "1.50.0"  # {x-release-please-version}
+__version__ = "1.49.0"  # {x-release-please-version}
diff --git a/google/cloud/aiplatform/v1/schema/predict/instance_v1/gapic_version.py b/google/cloud/aiplatform/v1/schema/predict/instance_v1/gapic_version.py
@@ -13,4 +13,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-__version__ = "1.50.0"  # {x-release-please-version}
+__version__ = "1.49.0"  # {x-release-please-version}
diff --git a/google/cloud/aiplatform/v1/schema/predict/params/gapic_version.py b/google/cloud/aiplatform/v1/schema/predict/params/gapic_version.py
@@ -13,4 +13,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-__version__ = "1.50.0"  # {x-release-please-version}
+__version__ = "1.49.0"  # {x-release-please-version}
diff --git a/google/cloud/aiplatform/v1/schema/predict/params_v1/gapic_version.py b/google/cloud/aiplatform/v1/schema/predict/params_v1/gapic_version.py
@@ -13,4 +13,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-__version__ = "1.50.0"  # {x-release-please-version}
+__version__ = "1.49.0"  # {x-release-please-version}
diff --git a/google/cloud/aiplatform/v1/schema/predict/prediction/gapic_version.py b/google/cloud/aiplatform/v1/schema/predict/prediction/gapic_version.py
@@ -13,4 +13,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-__version__ = "1.50.0"  # {x-release-please-version}
+__version__ = "1.49.0"  # {x-release-please-version}
diff --git a/google/cloud/aiplatform/v1/schema/predict/prediction_v1/gapic_version.py b/google/cloud/aiplatform/v1/schema/predict/prediction_v1/gapic_version.py
@@ -13,4 +13,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-__version__ = "1.50.0"  # {x-release-please-version}
+__version__ = "1.49.0"  # {x-release-please-version}
diff --git a/google/cloud/aiplatform/v1/schema/trainingjob/definition/gapic_version.py b/google/cloud/aiplatform/v1/schema/trainingjob/definition/gapic_version.py
@@ -13,4 +13,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-__version__ = "1.50.0"  # {x-release-please-version}
+__version__ = "1.49.0"  # {x-release-please-version}
diff --git a/google/cloud/aiplatform/v1/schema/trainingjob/definition_v1/gapic_version.py b/google/cloud/aiplatform/v1/schema/trainingjob/definition_v1/gapic_version.py
@@ -13,4 +13,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-__version__ = "1.50.0"  # {x-release-please-version}
+__version__ = "1.49.0"  # {x-release-please-version}
diff --git a/google/cloud/aiplatform/v1beta1/schema/predict/instance/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/predict/instance/gapic_version.py
@@ -13,4 +13,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-__version__ = "1.50.0"  # {x-release-please-version}
+__version__ = "1.49.0"  # {x-release-please-version}
diff --git a/google/cloud/aiplatform/v1beta1/schema/predict/instance_v1beta1/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/predict/instance_v1beta1/gapic_version.py
@@ -13,4 +13,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-__version__ = "1.50.0"  # {x-release-please-version}
+__version__ = "1.49.0"  # {x-release-please-version}
diff --git a/google/cloud/aiplatform/v1beta1/schema/predict/params/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/predict/params/gapic_version.py
@@ -13,4 +13,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-__version__ = "1.50.0"  # {x-release-please-version}
+__version__ = "1.49.0"  # {x-release-please-version}
diff --git a/google/cloud/aiplatform/v1beta1/schema/predict/params_v1beta1/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/predict/params_v1beta1/gapic_version.py
@@ -13,4 +13,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-__version__ = "1.50.0"  # {x-release-please-version}
+__version__ = "1.49.0"  # {x-release-please-version}
diff --git a/google/cloud/aiplatform/v1beta1/schema/predict/prediction/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/predict/prediction/gapic_version.py
@@ -13,4 +13,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-__version__ = "1.50.0"  # {x-release-please-version}
+__version__ = "1.49.0"  # {x-release-please-version}
diff --git a/google/cloud/aiplatform/v1beta1/schema/predict/prediction_v1beta1/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/predict/prediction_v1beta1/gapic_version.py
@@ -13,4 +13,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-__version__ = "1.50.0"  # {x-release-please-version}
+__version__ = "1.49.0"  # {x-release-please-version}
diff --git a/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition/gapic_version.py
@@ -13,4 +13,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-__version__ = "1.50.0"  # {x-release-please-version}
+__version__ = "1.49.0"  # {x-release-please-version}
diff --git a/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition_v1beta1/gapic_version.py b/google/cloud/aiplatform/v1beta1/schema/trainingjob/definition_v1beta1/gapic_version.py
@@ -13,4 +13,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-__version__ = "1.50.0"  # {x-release-please-version}
+__version__ = "1.49.0"  # {x-release-please-version}
diff --git a/google/cloud/aiplatform/version.py b/google/cloud/aiplatform/version.py
@@ -15,4 +15,4 @@
 # limitations under the License.
 #
 
-__version__ = "1.50.0"
+__version__ = "1.49.0"
diff --git a/google/cloud/aiplatform_v1/gapic_version.py b/google/cloud/aiplatform_v1/gapic_version.py
@@ -13,4 +13,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-__version__ = "1.50.0"  # {x-release-please-version}
+__version__ = "1.49.0"  # {x-release-please-version}
diff --git a/google/cloud/aiplatform_v1beta1/gapic_version.py b/google/cloud/aiplatform_v1beta1/gapic_version.py
@@ -13,4 +13,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-__version__ = "1.50.0"  # {x-release-please-version}
+__version__ = "1.49.0"  # {x-release-please-version}
diff --git a/pypi/_vertex_ai_placeholder/version.py b/pypi/_vertex_ai_placeholder/version.py
@@ -15,4 +15,4 @@
 # limitations under the License.
 #
 
-__version__ = "1.50.0"
+__version__ = "1.49.0"
diff --git a/samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1.json b/samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1.json
@@ -8,7 +8,7 @@
     ],
     "language": "PYTHON",
     "name": "google-cloud-aiplatform",
-    "version": "1.50.0"
+    "version": "1.49.0"
   },
   "snippets": [
     {
diff --git a/samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1beta1.json b/samples/generated_samples/snippet_metadata_google.cloud.aiplatform.v1beta1.json
@@ -8,7 +8,7 @@
     ],
     "language": "PYTHON",
     "name": "google-cloud-aiplatform",
-    "version": "1.50.0"
+    "version": "1.49.0"
   },
   "snippets": [
     {
diff --git a/setup.py b/setup.py
@@ -149,9 +149,9 @@
 ]
 
 langchain_extra_require = [
-    "langchain >= 0.1.13, < 0.2",
+    "langchain >= 0.1.16, < 0.2",
     "langchain-core < 0.2",
-    "langchain-google-vertexai < 0.2",
+    "langchain-google-vertexai < 2",
 ]
 
 langchain_testing_extra_require = list(
diff --git a/tests/unit/vertex_langchain/test_reasoning_engine_templates_langchain.py b/tests/unit/vertex_langchain/test_reasoning_engine_templates_langchain.py
@@ -13,7 +13,6 @@
 # limitations under the License.
 #
 import importlib
-import json
 from typing import Optional
 from unittest import mock
 
@@ -25,11 +24,9 @@
 import pytest
 
 
-from langchain_core import agents
-from langchain_core import messages
-from langchain_core import outputs
-from langchain_core import tools as lc_tools
+from langchain_core import prompts
 from langchain.load import dump as langchain_load_dump
+from langchain.agents.format_scratchpad import format_to_openai_function_messages
 from langchain.tools.base import StructuredTool
 
 
@@ -93,6 +90,18 @@ def setup_method(self):
             project=_TEST_PROJECT,
             location=_TEST_LOCATION,
         )
+        self.prompt = {
+            "input": lambda x: x["input"],
+            "agent_scratchpad": (
+                lambda x: format_to_openai_function_messages(x["intermediate_steps"])
+            ),
+        } | prompts.ChatPromptTemplate.from_messages(
+            [
+                ("user", "{input}"),
+                prompts.MessagesPlaceholder(variable_name="agent_scratchpad"),
+            ]
+        )
+        self.output_parser = mock.Mock()
 
     def teardown_method(self):
         initializer.global_pool.shutdown(wait=True)
@@ -105,24 +114,33 @@ def test_initialization(self):
         assert agent._runnable is None
 
     def test_initialization_with_tools(self):
+        tools = [
+            place_tool_query,
+            StructuredTool.from_function(place_photo_query),
+        ]
         agent = reasoning_engines.LangchainAgent(
             model=_TEST_MODEL,
-            tools=[
-                place_tool_query,
-                StructuredTool.from_function(place_photo_query),
-            ],
+            tools=tools,
         )
-        for tool in agent._tools:
-            assert isinstance(tool, lc_tools.BaseTool)
+        for tool, agent_tool in zip(tools, agent._tools):
+            assert isinstance(agent_tool, type(tool))
 
     def test_set_up(self, vertexai_init_mock):
-        agent = reasoning_engines.LangchainAgent(model=_TEST_MODEL)
+        agent = reasoning_engines.LangchainAgent(
+            model=_TEST_MODEL,
+            prompt=self.prompt,
+            output_parser=self.output_parser,
+        )
         assert agent._runnable is None
         agent.set_up()
         assert agent._runnable is not None
 
     def test_query(self, langchain_dump_mock):
-        agent = reasoning_engines.LangchainAgent(model=_TEST_MODEL)
+        agent = reasoning_engines.LangchainAgent(
+            model=_TEST_MODEL,
+            prompt=self.prompt,
+            output_parser=self.output_parser,
+        )
         agent._runnable = mock.Mock()
         mocks = mock.Mock()
         mocks.attach_mock(mock=agent._runnable, attribute="invoke")
@@ -132,63 +150,6 @@ def test_query(self, langchain_dump_mock):
         )
 
 
-class TestDefaultOutputParser:
-    def test_parse_result_function_call(self, vertexai_init_mock):
-        agent = reasoning_engines.LangchainAgent(model=_TEST_MODEL)
-        agent.set_up()
-        tool_input = {
-            "photo_reference": "abcd1234",
-            "maxwidth": _DEFAULT_PLACE_PHOTO_MAXWIDTH,
-        }
-        result = agent._output_parser.parse_result(
-            [
-                outputs.ChatGeneration(
-                    message=messages.AIMessage(
-                        content="",
-                        additional_kwargs={
-                            "function_call": {
-                                "name": "place_tool_query",
-                                "arguments": json.dumps(tool_input),
-                            },
-                        },
-                    )
-                )
-            ]
-        )
-        assert isinstance(result, agents.AgentActionMessageLog)
-        assert result.tool == "place_tool_query"
-        assert result.tool_input == tool_input
-
-    def test_parse_result_not_function_call(self, vertexai_init_mock):
-        agent = reasoning_engines.LangchainAgent(model=_TEST_MODEL)
-        agent.set_up()
-        content = "test content"
-        result = agent._output_parser.parse_result(
-            [
-                outputs.ChatGeneration(
-                    message=messages.AIMessage(content=content),
-                )
-            ]
-        )
-        assert isinstance(result, agents.AgentFinish)
-        assert result.return_values == {"output": content}
-        assert result.log == content
-
-
-class TestDefaultOutputParserErrors:
-    def test_parse_result_non_chat_generation_errors(self, vertexai_init_mock):
-        agent = reasoning_engines.LangchainAgent(model=_TEST_MODEL)
-        agent.set_up()
-        with pytest.raises(ValueError, match=r"only works on ChatGeneration"):
-            agent._output_parser.parse_result(["text"])
-
-    def test_parse_text_errors(self, vertexai_init_mock):
-        agent = reasoning_engines.LangchainAgent(model=_TEST_MODEL)
-        agent.set_up()
-        with pytest.raises(ValueError, match=r"Can only parse messages"):
-            agent._output_parser.parse("text")
-
-
 class TestConvertToolsOrRaise:
     def test_convert_tools_or_raise(self, vertexai_init_mock):
         pass
diff --git a/vertexai/preview/reasoning_engines/templates/langchain.py b/vertexai/preview/reasoning_engines/templates/langchain.py

Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,3 @@`
`1`	`1`	`{`
`2`		`- ".": "1.50.0"`
	`2`	`+ ".": "1.49.0"`
`3`	`3`	`}`
Original file line number	Diff line number	Diff line change
`@@ -13,4 +13,4 @@`
`13`	`13`	`# See the License for the specific language governing permissions and`
`14`	`14`	`# limitations under the License.`
`15`	`15`	`#`
`16`		`-__version__ = "1.50.0" # {x-release-please-version}`
	`16`	`+__version__ = "1.49.0" # {x-release-please-version}`
Original file line number	Diff line number	Diff line change
`@@ -15,4 +15,4 @@`
`15`	`15`	`# limitations under the License.`
`16`	`16`	`#`
`17`	`17`
`18`		`-__version__ = "1.50.0"`
	`18`	`+__version__ = "1.49.0"`
Original file line number	Diff line number	Diff line change
`@@ -8,7 +8,7 @@`
`8`	`8`	`],`
`9`	`9`	`"language": "PYTHON",`
`10`	`10`	`"name": "google-cloud-aiplatform",`
`11`		`- "version": "1.50.0"`
	`11`	`+ "version": "1.49.0"`
`12`	`12`	`},`
`13`	`13`	`"snippets": [`
`14`	`14`	`{`
Original file line number	Diff line number	Diff line change
`@@ -149,9 +149,9 @@`
`149`	`149`	`]`
`150`	`150`
`151`	`151`	`langchain_extra_require = [`
`152`		`- "langchain >= 0.1.13, < 0.2",`
	`152`	`+ "langchain >= 0.1.16, < 0.2",`
`153`	`153`	`"langchain-core < 0.2",`
`154`		`- "langchain-google-vertexai < 0.2",`
	`154`	`+ "langchain-google-vertexai < 2",`
`155`	`155`	`]`
`156`	`156`
`157`	`157`	`langchain_testing_extra_require = list(`