fix: GenAI SDK client - handle os paths for .whl in requirements when deploying to agent engine

PiperOrigin-RevId: 788584767

Author: yeesian

tests/unit/vertex_langchain/test_agent_engines.py

@@ -3251,6 +3251,19 @@ def test_invalid_requirement_warning(self, caplog):
  _utils.parse_constraints(["invalid requirement line"])
  assert "Failed to parse constraint" in caplog.text
 
+ def test_requirements_with_whl_files(self):
+ whl_files = [
+ "wxPython-4.2.4-cp39-cp39-macosx_12_0_x86_64.whl",
+ "/content/wxPython-4.2.3-cp39-cp39-macosx_12_0_x86_64.whl",
+ "https://wxpython.org/Phoenix/snapshot-builds/wxPython-4.2.2-cp38-cp38-macosx_12_0_x86_64.whl",
+ ]
+ result = _utils.parse_constraints(whl_files)
+ assert result == {
+ "wxPython-4.2.2-cp38-cp38-macosx_12_0_x86_64.whl": None,
+ "wxPython-4.2.3-cp39-cp39-macosx_12_0_x86_64.whl": None,
+ "wxPython-4.2.4-cp39-cp39-macosx_12_0_x86_64.whl": None,
+ }
+
  def test_compare_requirements_with_required_packages(self):
  requirements = {"requests": "2.0.0"}
  constraints = ["requests==1.0.0"]

tests/unit/vertexai/genai/test_agent_engines.py

@@ -730,6 +730,14 @@ def set_up(self):
  pass
 
 
+def _create_fake_object_with_module(module_name):
+ class FakeObject:
+ pass
+
+ FakeObject.__module__ = module_name
+ return FakeObject()
+
+
 @pytest.mark.usefixtures("google_auth_mock")
 class TestAgentEngineHelpers:
  def setup_method(self):
@@ -919,6 +927,102 @@ def test_register_api_methods(self):
  "description"
  )
 
+ @pytest.mark.usefixtures("caplog")
+ def test_invalid_requirement_warning(self, caplog):
+ _agent_engines_utils._parse_constraints(["invalid requirement line"])
+ assert "Failed to parse constraint" in caplog.text
+
+ def test_requirements_with_whl_files(self):
+ whl_files = [
+ "wxPython-4.2.4-cp39-cp39-macosx_12_0_x86_64.whl",
+ "/content/wxPython-4.2.3-cp39-cp39-macosx_12_0_x86_64.whl",
+ "https://wxpython.org/Phoenix/snapshot-builds/wxPython-4.2.2-cp38-cp38-macosx_12_0_x86_64.whl",
+ ]
+ result = _agent_engines_utils._parse_constraints(whl_files)
+ assert result == {
+ "wxPython-4.2.2-cp38-cp38-macosx_12_0_x86_64.whl": None,
+ "wxPython-4.2.3-cp39-cp39-macosx_12_0_x86_64.whl": None,
+ "wxPython-4.2.4-cp39-cp39-macosx_12_0_x86_64.whl": None,
+ }
+
+ def test_compare_requirements_with_required_packages(self):
+ requirements = {"requests": "2.0.0"}
+ constraints = ["requests==1.0.0"]
+ result = _agent_engines_utils._compare_requirements(requirements, constraints)
+ assert result == {
+ "actions": {"append": set()},
+ "warnings": {
+ "incompatible": {"requests==2.0.0 (required: ==1.0.0)"},
+ "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 = _agent_engines_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 = _agent_engines_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 = _agent_engines_utils._BASE_MODULES
+ monkeypatch.setattr(
+ _agent_engines_utils,
+ "_BASE_MODULES",
+ set(original_base) | {"requests"},
+ )
+ requirements = _agent_engines_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 = _agent_engines_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",
+ }
+
 
 @pytest.mark.usefixtures("google_auth_mock")
 class TestAgentEngine:

vertexai/_genai/_agent_engines_utils.py

@@ -686,6 +686,8 @@ def _parse_constraints(
  result: Dict[str, Optional[_SpecifierSet]] = {}
  for constraint in constraints:
  try:
+ if constraint.endswith(".whl"):
+ constraint = os.path.basename(constraint)
  requirement = requirements.Requirement(constraint)
  except Exception as e:
  logger.warning(f"Failed to parse constraint: {constraint}. Exception: {e}")

vertexai/agent_engines/_utils.py

@@ -16,6 +16,7 @@
 import dataclasses
 import inspect
 import json
+import os
 import sys
 import types
 import typing
@@ -359,6 +360,8 @@ def parse_constraints(
  result = {}
  for constraint in constraints:
  try:
+ if constraint.endswith(".whl"):
+ constraint = os.path.basename(constraint)
  requirement = requirements.Requirement(constraint)
  except Exception as e:
  LOGGER.warning(f"Failed to parse constraint: {constraint}. Exception: {e}")