From 16e0af92a419339b4292f83856ef162aeb548875 Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Tue, 28 Apr 2026 13:59:27 -0700 Subject: [PATCH] feat: GenAI SDK client - Adding client-based SDK for Skill Registry - GET method PiperOrigin-RevId: 907166269 --- .../unit/vertexai/genai/test_genai_skills.py | 48 +++++++++++++++++++ vertexai/_genai/client.py | 21 ++++++++ 2 files changed, 69 insertions(+) create mode 100644 tests/unit/vertexai/genai/test_genai_skills.py diff --git a/tests/unit/vertexai/genai/test_genai_skills.py b/tests/unit/vertexai/genai/test_genai_skills.py new file mode 100644 index 0000000000..68e29b1e16 --- /dev/null +++ b/tests/unit/vertexai/genai/test_genai_skills.py @@ -0,0 +1,48 @@ +# //third_party/py/google/cloud/aiplatform/tests/unit/vertexai/genai/test_genai_skills.py +import pytest +from unittest import mock +from vertexai import client as vertexai_client +from vertexai import _genai as genai +from google.genai import client as genai_client +from google.genai import deps + +@pytest.fixture +def mock_genai_client(): + return mock.create_autospec(genai_client.Client) + +@pytest.fixture +def skills_client(mock_genai_client): + creds = mock.MagicMock() + creds.token = "test_token" + client = vertexai_client.Client( + project="test-project", location="test-location", credentials=creds + ) + client._genai_client = mock_genai_client + return client.skills + +class TestGenaiSkills: + mock_get_skill_response = { + "name": "projects/test-project/locations/test-location/skills/test-skill", + "displayName": "My Test Skill", + # Add other expected fields from the Skill proto + } + + def test_get_skill(self, skills_client, mock_genai_client): + """Tests the get_skill method.""" + mock_genai_client.post.return_value = deps.Response( + result=self.mock_get_skill_response, + request=mock.MagicMock(), + response=mock.MagicMock(), + ) + + skill_name = "projects/test-project/locations/test-location/skills/test-skill" + skill = skills_client.get(name=skill_name) + + mock_genai_client.post.assert_called_once() + call_args = mock_genai_client.post.call_args + assert call_args[0][0] == skill_name + assert call_args[1]["method"] == "GET" + + assert isinstance(skill, genai.types.Skill) + assert skill.name == skill_name + assert skill.display_name == "My Test Skill" diff --git a/vertexai/_genai/client.py b/vertexai/_genai/client.py index 2e43782554..77d8a9aadc 100644 --- a/vertexai/_genai/client.py +++ b/vertexai/_genai/client.py @@ -35,6 +35,7 @@ prompt_optimizer as prompt_optimizer_module, ) from vertexai._genai import prompts as prompts_module + from vertexai._genai import skills as skills_module from vertexai._genai import live as live_module @@ -52,6 +53,7 @@ def __init__(self, api_client: genai_client.BaseApiClient): # type: ignore[name self._prompt_optimizer: Optional[ModuleType] = None self._prompts: Optional[ModuleType] = None self._datasets: Optional[ModuleType] = None + self._skills: Optional[ModuleType] = None @property @_common.experimental_warning( @@ -124,6 +126,15 @@ def datasets(self) -> "datasets_module.AsyncDatasets": ) return self._datasets.AsyncDatasets(self._api_client) # type: ignore[no-any-return] + @property + def skills(self) -> "skills_module.AsyncSkills": + if self._skills is None: + self._skills = importlib.import_module( + ".skills", + __package__, + ) + return self._skills.AsyncSkills(self._api_client) # type: ignore[no-any-return] + async def aclose(self) -> None: """Closes the async client explicitly. @@ -239,6 +250,7 @@ def __init__( self._agent_engines: Optional[ModuleType] = None self._prompts: Optional[ModuleType] = None self._datasets: Optional[ModuleType] = None + self._skills: Optional[ModuleType] = None @property def evals(self) -> "evals_module.Evals": @@ -335,3 +347,12 @@ def datasets(self) -> "datasets_module.Datasets": __package__, ) return self._datasets.Datasets(self._api_client) # type: ignore[no-any-return] + + @property + def skills(self) -> "skills_module.Skills": + if self._skills is None: + self._skills = importlib.import_module( + ".skills", + __package__, + ) + return self._skills.Skills(self._api_client) # type: ignore[no-any-return]