feat: Add pydantic to default required packages for agent engines

shawn-yang-google · copybara-github · commit 6960ae4fb2e3 · 2025-04-21T11:35:08.000-07:00
PiperOrigin-RevId: 749861402
diff --git a/tests/unit/vertex_langchain/test_agent_engines.py b/tests/unit/vertex_langchain/test_agent_engines.py
@@ -399,6 +399,11 @@ def register_operations(self) -> Dict[str, List[str]]:
 _TEST_STREAM_QUERY_SCHEMAS = [
     _TEST_AGENT_ENGINE_STREAM_QUERY_SCHEMA,
 ]
+_TEST_PACKAGE_DISTRIBUTIONS = {
+    "requests": ["requests"],
+    "cloudpickle": ["cloudpickle"],
+    "pydantic": ["pydantic"],
+}
 
 
 def _create_empty_fake_package(package_name: str) -> str:
@@ -507,6 +512,16 @@ def importlib_metadata_version_mock():
     with mock.patch.object(
         importlib.metadata, "version"
     ) as importlib_metadata_version_mock:
+
+        def get_version(pkg):
+            versions = {
+                "requests": "2.0.0",
+                "cloudpickle": "3.0.0",
+                "pydantic": "1.11.1",
+            }
+            return versions.get(pkg, "unknown")
+
+        importlib_metadata_version_mock.side_effect = get_version
         yield importlib_metadata_version_mock
 
 
@@ -616,6 +631,14 @@ def unregister_api_methods_mock():
         yield unregister_api_methods_mock
 
 
+def create_fake_object_with_module(module_name):
+    class FakeObject:
+        pass
+
+    FakeObject.__module__ = module_name
+    return FakeObject()
+
+
 class InvalidCapitalizeEngineWithoutQuerySelf:
     """A sample Agent Engine with an invalid query method."""
 
@@ -2519,3 +2542,69 @@ def test_compare_requirements_with_required_packages(self):
                 "missing": set(),
             },
         }
+
+    @pytest.mark.usefixtures("importlib_metadata_version_mock")
+    def test_scan_simple_object(self):
+        """Test scanning an object importing a known third-party package."""
+        fake_obj = create_fake_object_with_module("requests")
+        requirements = _utils.scan_requirements(
+            fake_obj,
+            package_distributions=_TEST_PACKAGE_DISTRIBUTIONS,
+        )
+        assert requirements == {
+            "cloudpickle": "3.0.0",
+            "pydantic": "1.11.1",
+            "requests": "2.0.0",
+        }
+
+    @pytest.mark.usefixtures("importlib_metadata_version_mock")
+    def test_scan_object_with_stdlib_module(self):
+        """Test that stdlib modules are ignored by default."""
+        fake_obj_stdlib = create_fake_object_with_module("json")
+        requirements = _utils.scan_requirements(
+            fake_obj_stdlib,
+            package_distributions=_TEST_PACKAGE_DISTRIBUTIONS,
+        )
+        # Requirements should not contain 'json',
+        # because 'json' is a stdlib module.
+        assert requirements == {
+            "cloudpickle": "3.0.0",
+            "pydantic": "1.11.1",
+        }
+
+    @pytest.mark.usefixtures("importlib_metadata_version_mock")
+    def test_scan_with_default_ignore_modules(self, monkeypatch):
+        """Test implicitly ignoring a module."""
+        fake_obj = create_fake_object_with_module("requests")
+        original_base = _utils._BASE_MODULES
+        monkeypatch.setattr(
+            _utils,
+            "_BASE_MODULES",
+            set(original_base) | {"requests"},
+        )
+        requirements = _utils.scan_requirements(
+            fake_obj,
+            package_distributions=_TEST_PACKAGE_DISTRIBUTIONS,
+        )
+        # Requirements should not contain 'requests',
+        # because 'requests' is implicitly ignored in `_BASE_MODULES`.
+        assert requirements == {
+            "cloudpickle": "3.0.0",
+            "pydantic": "1.11.1",
+        }
+
+    @pytest.mark.usefixtures("importlib_metadata_version_mock")
+    def test_scan_with_explicit_ignore_modules(self):
+        """Test explicitly ignoring a module."""
+        fake_obj = create_fake_object_with_module("requests")
+        requirements = _utils.scan_requirements(
+            fake_obj,
+            ignore_modules=["requests"],
+            package_distributions=_TEST_PACKAGE_DISTRIBUTIONS,
+        )
+        # Requirements should not contain 'requests',
+        # because 'requests' is explicitly ignored in `ignore_modules`.
+        assert requirements == {
+            "cloudpickle": "3.0.0",
+            "pydantic": "1.11.1",
+        }
diff --git a/vertexai/agent_engines/_utils.py b/vertexai/agent_engines/_utils.py
@@ -110,7 +110,7 @@ class _RequirementsValidationResult(TypedDict):
 LOGGER = base.Logger("vertexai.agent_engines")
 
 _BASE_MODULES = set(_BUILTIN_MODULE_NAMES + tuple(_STDLIB_MODULE_NAMES))
-_DEFAULT_REQUIRED_PACKAGES = frozenset(["cloudpickle"])
+_DEFAULT_REQUIRED_PACKAGES = frozenset(["cloudpickle", "pydantic"])
 _ACTIONS_KEY = "actions"
 _ACTION_APPEND = "append"
 _WARNINGS_KEY = "warnings"
