Skip to content

Commit 6f89fd3

Browse files
vertex-sdk-botcopybara-github
authored andcommitted
feat: Add client for Vertex AI Skill Registry - Get Skill method
PiperOrigin-RevId: 911401004
1 parent 062e254 commit 6f89fd3

6 files changed

Lines changed: 444 additions & 0 deletions

File tree

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
"""Tests the skills.get() method against the autopush endpoint."""
2+
3+
from google.api_core import exceptions
4+
from tests.unit.vertexai.genai.replays import pytest_helper
5+
import pytest
6+
7+
PROJECT_ID = "demo-project"
8+
REGION = "us-central1"
9+
SKILL_ID = "7184367305562783744"
10+
# target the autopush sandbox endpoint for the Skill Registry API
11+
ENDPOINT = f"{REGION}-autopush-aiplatform.sandbox.googleapis.com"
12+
13+
14+
pytestmark = pytest_helper.setup(
15+
file=__file__,
16+
globals_for_file=globals(),
17+
)
18+
19+
20+
def test_get_skill(client): # client fixture is injected by pytest_helper.setup
21+
"""Tests the skills.get() method against the autopush endpoint."""
22+
23+
client._api_client._http_options.base_url = (
24+
"https://us-central1-autopush-aiplatform.sandbox.googleapis.com"
25+
)
26+
skill_name = f"projects/{PROJECT_ID}/locations/{REGION}/skills/{SKILL_ID}"
27+
28+
try:
29+
skill = client.skills.get(name=skill_name)
30+
assert skill.name == skill_name
31+
32+
except exceptions.GoogleAPIError as e:
33+
pytest.fail(f"Error calling client.skills.get(): {e}")
Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
# //third_party/py/google/cloud/aiplatform/tests/unit/vertexai/genai/test_genai_skills.py
2+
import json
3+
from unittest import mock
4+
from vertexai import _genai as genai
5+
from vertexai._genai import client as vertexai_client
6+
from google.genai import types as genai_types
7+
import pytest
8+
9+
10+
@pytest.fixture
11+
def skills_client():
12+
creds = mock.MagicMock()
13+
creds.token = "test_token"
14+
client = vertexai_client.Client(
15+
project="test-project", location="test-location", credentials=creds
16+
)
17+
return client.skills
18+
19+
20+
class TestGenaiSkills:
21+
mock_get_skill_response = {
22+
"name": "projects/test-project/locations/test-location/skills/test-skill",
23+
"displayName": "My Test Skill",
24+
}
25+
26+
def test_get_skill(self, skills_client):
27+
"""Tests the get_skill method."""
28+
with mock.patch.object(skills_client._api_client, "request") as request_mock:
29+
request_mock.return_value = genai_types.HttpResponse(
30+
body=json.dumps(self.mock_get_skill_response)
31+
)
32+
skill_name = (
33+
"projects/test-project/locations/test-location/skills/test-skill"
34+
)
35+
skill = skills_client.get(name=skill_name)
36+
request_mock.assert_called_with(
37+
"get",
38+
skill_name,
39+
{"_url": {"name": skill_name}},
40+
None,
41+
)
42+
assert isinstance(skill, genai.types.Skill)
43+
assert skill.name == skill_name
44+
assert skill.display_name == "My Test Skill"

vertexai/_genai/client.py

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
prompt_optimizer as prompt_optimizer_module,
3636
)
3737
from vertexai._genai import prompts as prompts_module
38+
from vertexai._genai import skills as skills_module
3839
from vertexai._genai import live as live_module
3940

4041

@@ -52,6 +53,7 @@ def __init__(self, api_client: genai_client.BaseApiClient): # type: ignore[name
5253
self._prompt_optimizer: Optional[ModuleType] = None
5354
self._prompts: Optional[ModuleType] = None
5455
self._datasets: Optional[ModuleType] = None
56+
self._skills: Optional[ModuleType] = None
5557

5658
@property
5759
@_common.experimental_warning(
@@ -124,6 +126,15 @@ def datasets(self) -> "datasets_module.AsyncDatasets":
124126
)
125127
return self._datasets.AsyncDatasets(self._api_client) # type: ignore[no-any-return]
126128

129+
@property
130+
def skills(self) -> "skills_module.AsyncSkills":
131+
if self._skills is None:
132+
self._skills = importlib.import_module(
133+
".skills",
134+
__package__,
135+
)
136+
return self._skills.AsyncSkills(self._api_client) # type: ignore[no-any-return]
137+
127138
async def aclose(self) -> None:
128139
"""Closes the async client explicitly.
129140
@@ -239,6 +250,7 @@ def __init__(
239250
self._agent_engines: Optional[ModuleType] = None
240251
self._prompts: Optional[ModuleType] = None
241252
self._datasets: Optional[ModuleType] = None
253+
self._skills: Optional[ModuleType] = None
242254

243255
@property
244256
def evals(self) -> "evals_module.Evals":
@@ -335,3 +347,12 @@ def datasets(self) -> "datasets_module.Datasets":
335347
__package__,
336348
)
337349
return self._datasets.Datasets(self._api_client) # type: ignore[no-any-return]
350+
351+
@property
352+
def skills(self) -> "skills_module.Skills":
353+
if self._skills is None:
354+
self._skills = importlib.import_module(
355+
".skills",
356+
__package__,
357+
)
358+
return self._skills.Skills(self._api_client) # type: ignore[no-any-return]

vertexai/_genai/skills.py

Lines changed: 192 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,192 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
#
15+
16+
# Code generated by the Google Gen AI SDK generator DO NOT EDIT.
17+
18+
import json
19+
import logging
20+
from typing import Any, Optional, Union
21+
from urllib.parse import urlencode
22+
23+
from google.genai import _api_module
24+
from google.genai import _common
25+
from google.genai._common import get_value_by_path as getv
26+
from google.genai._common import set_value_by_path as setv
27+
28+
from . import types
29+
30+
logger = logging.getLogger("vertexai_genai.skills")
31+
32+
33+
def _GetSkillRequestParameters_to_vertex(
34+
from_object: Union[dict[str, Any], object],
35+
parent_object: Optional[dict[str, Any]] = None,
36+
) -> dict[str, Any]:
37+
to_object: dict[str, Any] = {}
38+
if getv(from_object, ["name"]) is not None:
39+
setv(to_object, ["_url", "name"], getv(from_object, ["name"]))
40+
41+
if getv(from_object, ["config"]) is not None:
42+
setv(to_object, ["config"], getv(from_object, ["config"]))
43+
44+
return to_object
45+
46+
47+
class Skills(_api_module.BaseModule):
48+
"""Class for managing Skills in the Skill Registry."""
49+
50+
def get(
51+
self, *, name: str, config: Optional[types.GetSkillConfigOrDict] = None
52+
) -> types.Skill:
53+
"""
54+
Gets a Skill.
55+
"""
56+
57+
parameter_model = types._GetSkillRequestParameters(
58+
name=name,
59+
config=config,
60+
)
61+
62+
request_url_dict: Optional[dict[str, str]]
63+
if not self._api_client.vertexai:
64+
raise ValueError(
65+
"This method is only supported in the Gemini Enterprise Agent Platform (previously known as Vertex AI) client."
66+
)
67+
else:
68+
request_dict = _GetSkillRequestParameters_to_vertex(parameter_model)
69+
request_url_dict = request_dict.get("_url")
70+
if request_url_dict:
71+
path = "{name}".format_map(request_url_dict)
72+
else:
73+
path = "{name}"
74+
75+
query_params = request_dict.get("_query")
76+
if query_params:
77+
path = f"{path}?{urlencode(query_params)}"
78+
# TODO: remove the hack that pops config.
79+
request_dict.pop("config", None)
80+
81+
http_options: Optional[types.HttpOptions] = None
82+
if (
83+
parameter_model.config is not None
84+
and parameter_model.config.http_options is not None
85+
):
86+
http_options = parameter_model.config.http_options
87+
88+
request_dict = _common.convert_to_dict(request_dict)
89+
request_dict = _common.encode_unserializable_types(request_dict)
90+
91+
response = self._api_client.request("get", path, request_dict, http_options)
92+
93+
response_dict = {} if not response.body else json.loads(response.body)
94+
95+
return_value = types.Skill._from_response(
96+
response=response_dict,
97+
kwargs=(
98+
{
99+
"config": {
100+
"response_schema": getattr(
101+
parameter_model.config, "response_schema", None
102+
),
103+
"response_json_schema": getattr(
104+
parameter_model.config, "response_json_schema", None
105+
),
106+
"include_all_fields": getattr(
107+
parameter_model.config, "include_all_fields", None
108+
),
109+
}
110+
}
111+
if getattr(parameter_model, "config", None)
112+
else {}
113+
),
114+
)
115+
116+
self._api_client._verify_response(return_value)
117+
return return_value
118+
119+
120+
class AsyncSkills(_api_module.BaseModule):
121+
"""Class for managing Skills in the Skill Registry."""
122+
123+
async def get(
124+
self, *, name: str, config: Optional[types.GetSkillConfigOrDict] = None
125+
) -> types.Skill:
126+
"""
127+
Gets a Skill.
128+
"""
129+
130+
parameter_model = types._GetSkillRequestParameters(
131+
name=name,
132+
config=config,
133+
)
134+
135+
request_url_dict: Optional[dict[str, str]]
136+
if not self._api_client.vertexai:
137+
raise ValueError(
138+
"This method is only supported in the Gemini Enterprise Agent Platform (previously known as Vertex AI) client."
139+
)
140+
else:
141+
request_dict = _GetSkillRequestParameters_to_vertex(parameter_model)
142+
request_url_dict = request_dict.get("_url")
143+
if request_url_dict:
144+
path = "{name}".format_map(request_url_dict)
145+
else:
146+
path = "{name}"
147+
148+
query_params = request_dict.get("_query")
149+
if query_params:
150+
path = f"{path}?{urlencode(query_params)}"
151+
# TODO: remove the hack that pops config.
152+
request_dict.pop("config", None)
153+
154+
http_options: Optional[types.HttpOptions] = None
155+
if (
156+
parameter_model.config is not None
157+
and parameter_model.config.http_options is not None
158+
):
159+
http_options = parameter_model.config.http_options
160+
161+
request_dict = _common.convert_to_dict(request_dict)
162+
request_dict = _common.encode_unserializable_types(request_dict)
163+
164+
response = await self._api_client.async_request(
165+
"get", path, request_dict, http_options
166+
)
167+
168+
response_dict = {} if not response.body else json.loads(response.body)
169+
170+
return_value = types.Skill._from_response(
171+
response=response_dict,
172+
kwargs=(
173+
{
174+
"config": {
175+
"response_schema": getattr(
176+
parameter_model.config, "response_schema", None
177+
),
178+
"response_json_schema": getattr(
179+
parameter_model.config, "response_json_schema", None
180+
),
181+
"include_all_fields": getattr(
182+
parameter_model.config, "include_all_fields", None
183+
),
184+
}
185+
}
186+
if getattr(parameter_model, "config", None)
187+
else {}
188+
),
189+
)
190+
191+
self._api_client._verify_response(return_value)
192+
return return_value

vertexai/_genai/types/__init__.py

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,7 @@
9090
from .common import _GetSandboxEnvironmentSnapshotRequestParameters
9191
from .common import _GetSandboxEnvironmentTemplateOperationParameters
9292
from .common import _GetSandboxEnvironmentTemplateRequestParameters
93+
from .common import _GetSkillRequestParameters
9394
from .common import _IngestEventsRequestParameters
9495
from .common import _ListAgentEngineMemoryRequestParameters
9596
from .common import _ListAgentEngineMemoryRevisionsRequestParameters
@@ -608,6 +609,9 @@
608609
from .common import GetSandboxEnvironmentTemplateConfig
609610
from .common import GetSandboxEnvironmentTemplateConfigDict
610611
from .common import GetSandboxEnvironmentTemplateConfigOrDict
612+
from .common import GetSkillConfig
613+
from .common import GetSkillConfigDict
614+
from .common import GetSkillConfigOrDict
611615
from .common import IdentityType
612616
from .common import Importance
613617
from .common import IngestEventsConfig
@@ -1243,6 +1247,10 @@
12431247
from .common import SessionEventDict
12441248
from .common import SessionEventOrDict
12451249
from .common import SessionOrDict
1250+
from .common import Skill
1251+
from .common import SkillDict
1252+
from .common import SkillOrDict
1253+
from .common import SkillState
12461254
from .common import State
12471255
from .common import Strategy
12481256
from .common import StructuredMemoryConfig
@@ -2478,6 +2486,12 @@
24782486
"UpdatePromptConfig",
24792487
"UpdatePromptConfigDict",
24802488
"UpdatePromptConfigOrDict",
2489+
"GetSkillConfig",
2490+
"GetSkillConfigDict",
2491+
"GetSkillConfigOrDict",
2492+
"Skill",
2493+
"SkillDict",
2494+
"SkillOrDict",
24812495
"PromptOptimizerConfig",
24822496
"PromptOptimizerConfigDict",
24832497
"PromptOptimizerConfigOrDict",
@@ -2591,6 +2605,7 @@
25912605
"OptimizeTarget",
25922606
"MemoryMetadataMergeStrategy",
25932607
"GenerateMemoriesResponseGeneratedMemoryAction",
2608+
"SkillState",
25942609
"PromptOptimizerMethod",
25952610
"OptimizationMethod",
25962611
"PromptData",
@@ -2714,6 +2729,7 @@
27142729
"_CustomJobParameters",
27152730
"_GetCustomJobParameters",
27162731
"_OptimizeRequestParameters",
2732+
"_GetSkillRequestParameters",
27172733
"evals",
27182734
"agent_engines",
27192735
"prompts",

0 commit comments

Comments
 (0)