chore: Add unit tests for langchain template.

yeesian · copybara-github · commit 894c73f2e971 · 2024-04-12T11:15:00.000-07:00
PiperOrigin-RevId: 624236094
diff --git a/noxfile.py b/noxfile.py
@@ -68,6 +68,7 @@
 nox.options.sessions = [
     "unit",
     "unit_ray",
+    "unit_langchain",
     "system",
     "cover",
     "lint",
@@ -181,6 +182,7 @@ def default(session):
         "--cov-report=",
         "--cov-fail-under=0",
         "--ignore=tests/unit/vertex_ray",
+        "--ignore=tests/unit/vertex_langchain",
         os.path.join("tests", "unit"),
         *session.posargs,
     )
@@ -219,6 +221,32 @@ def unit_ray(session, ray):
     )
 
 
+@nox.session(python=UNIT_TEST_PYTHON_VERSIONS)
+def unit_langchain(session):
+    # Install all test dependencies, then install this package in-place.
+
+    constraints_path = str(CURRENT_DIRECTORY / "testing" / f"constraints-langchain.txt")
+    standard_deps = UNIT_TEST_STANDARD_DEPENDENCIES + UNIT_TEST_DEPENDENCIES
+    session.install(*standard_deps, "-c", constraints_path)
+
+    # Install langchain extras
+    session.install("-e", ".[langchain_testing]", "-c", constraints_path)
+
+    # Run py.test against the unit tests.
+    session.run(
+        "py.test",
+        "--quiet",
+        f"--junitxml=unit_langchain_sponge_log.xml",
+        "--cov=google",
+        "--cov-append",
+        "--cov-config=.coveragerc",
+        "--cov-report=",
+        "--cov-fail-under=0",
+        os.path.join("tests", "unit", "vertex_langchain"),
+        *session.posargs,
+    )
+
+
 def install_systemtest_dependencies(session, *constraints):
     # Use pre-release gRPC for system tests.
     # Exclude version 1.52.0rc1 which has a known issue.
diff --git a/setup.py b/setup.py
@@ -260,6 +260,7 @@
         "reasoningengine": reasoning_engine_extra_require,
         "rapid_evaluation": rapid_evaluation_extra_require,
         "langchain": langchain_extra_require,
+        "langchain_testing": langchain_extra_require,
     },
     python_requires=">=3.8",
     classifiers=[
diff --git a/testing/constraints-langchain.txt b/testing/constraints-langchain.txt
@@ -0,0 +1,3 @@
+langchain==0.1.15
+langchain-core==0.1.40
+langchain-google-vertexai==0.1.2
diff --git a/tests/unit/vertex_langchain/test_reasoning_engine_templates_langchain.py b/tests/unit/vertex_langchain/test_reasoning_engine_templates_langchain.py
@@ -0,0 +1,180 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import importlib
+import json
+from typing import Optional
+from unittest import mock
+
+from google import auth
+from google.auth import credentials as auth_credentials
+import vertexai
+from google.cloud.aiplatform import initializer
+from vertexai.preview import reasoning_engines
+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
+
+
+_DEFAULT_PLACE_TOOL_ACTIVITY = "museums"
+_DEFAULT_PLACE_TOOL_PAGE_SIZE = 3
+_DEFAULT_PLACE_PHOTO_MAXWIDTH = 400
+_TEST_LOCATION = "us-central1"
+_TEST_PROJECT = "test-project"
+_TEST_MODEL = "gemini-1.0-pro"
+
+
+def place_tool_query(
+    city: str,
+    activity: str = _DEFAULT_PLACE_TOOL_ACTIVITY,
+    page_size: int = _DEFAULT_PLACE_TOOL_PAGE_SIZE,
+):
+    """Searches the city for recommendations on the activity."""
+    return {"city": city, "activity": activity, "page_size": page_size}
+
+
+def place_photo_query(
+    photo_reference: str,
+    maxwidth: int = _DEFAULT_PLACE_PHOTO_MAXWIDTH,
+    maxheight: Optional[int] = None,
+):
+    """Returns the photo for a given reference."""
+    result = {"photo_reference": photo_reference, "maxwidth": maxwidth}
+    if maxheight:
+        result["maxheight"] = maxheight
+    return result
+
+
+@pytest.fixture(scope="module")
+def google_auth_mock():
+    with mock.patch.object(auth, "default") as google_auth_mock:
+        google_auth_mock.return_value = (
+            auth_credentials.AnonymousCredentials(),
+            _TEST_PROJECT,
+        )
+        yield google_auth_mock
+
+
+@pytest.fixture
+def vertexai_init_mock():
+    with mock.patch.object(vertexai, "init") as vertexai_init_mock:
+        yield vertexai_init_mock
+
+
+@pytest.mark.usefixtures("google_auth_mock")
+class TestLangchainAgent:
+    def setup_method(self):
+        importlib.reload(initializer)
+        importlib.reload(vertexai)
+        vertexai.init(
+            project=_TEST_PROJECT,
+            location=_TEST_LOCATION,
+        )
+
+    def teardown_method(self):
+        initializer.global_pool.shutdown(wait=True)
+
+    def test_initialization(self):
+        agent = reasoning_engines.LangchainAgent(model=_TEST_MODEL)
+        assert agent._model_name == _TEST_MODEL
+        assert agent._project == _TEST_PROJECT
+        assert agent._location == _TEST_LOCATION
+        assert agent._runnable is None
+
+    def test_initialization_with_tools(self):
+        agent = reasoning_engines.LangchainAgent(
+            model=_TEST_MODEL,
+            tools=[
+                place_tool_query,
+                place_photo_query,
+            ],
+        )
+        for tool in agent._tools:
+            assert isinstance(tool, lc_tools.BaseTool)
+
+    def test_set_up(self, vertexai_init_mock):
+        agent = reasoning_engines.LangchainAgent(model=_TEST_MODEL)
+        assert agent._runnable is None
+        agent.set_up()
+        assert agent._runnable is not None
+
+    def test_query(self):
+        agent = reasoning_engines.LangchainAgent(model=_TEST_MODEL)
+        agent._runnable = mock.Mock()
+        mocks = mock.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)]
+        )
+
+
+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")

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+langchain==0.1.15`
	`2`	`+langchain-core==0.1.40`
	`3`	`+langchain-google-vertexai==0.1.2`