Skip to content

Commit 83a44ac

Browse files
vertex-sdk-botcopybara-github
authored andcommitted
feat: GenAI SDK client - Adding client-based SDK for Skill Registry - GET method
PiperOrigin-RevId: 907166269
1 parent f48e54b commit 83a44ac

6 files changed

Lines changed: 597 additions & 0 deletions

File tree

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
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 = "srbai-testing"
8+
REGION = "us-central1"
9+
# SKILL_ID = "5578834038405201920"
10+
SKILL_ID = "7184367305562783744"
11+
ENDPOINT = f"{REGION}-autopush-aiplatform.sandbox.googleapis.com"
12+
13+
# # Configure HTTP options to target the autopush endpoint
14+
# my_http_options = genai_types.HttpOptions(
15+
# api_version="v1beta1",
16+
# base_url=f"https://{ENDPOINT}/v1beta1/" # <---APPENDED /v1beta1/ here
17+
# )
18+
19+
pytestmark = pytest_helper.setup(
20+
file=__file__,
21+
globals_for_file=globals(),
22+
# http_options=my_http_options,
23+
)
24+
25+
26+
def test_get_skill(client): # client fixture is injected by pytest_helper.setup
27+
"""Tests the skills.get() method against the autopush endpoint."""
28+
29+
client._api_client._http_options.base_url = (
30+
"https://us-central1-autopush-aiplatform.sandbox.googleapis.com"
31+
)
32+
skill_name = f"projects/{PROJECT_ID}/locations/{REGION}/skills/{SKILL_ID}"
33+
34+
try:
35+
skill = client.skills.get(name=skill_name)
36+
assert skill.name == skill_name
37+
38+
except exceptions.GoogleAPIError as e:
39+
pytest.fail(f"Error calling client.skills.get(): {e}")
Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
# //third_party/py/google/cloud/aiplatform/tests/unit/vertexai/genai/test_genai_skills.py
2+
import pytest
3+
from unittest import mock
4+
from vertexai import client as vertexai_client
5+
from vertexai import _genai as genai
6+
from google.genai import client as genai_client
7+
from google.genai import deps
8+
9+
10+
@pytest.fixture
11+
def mock_genai_client():
12+
return mock.create_autospec(genai_client.Client)
13+
14+
15+
@pytest.fixture
16+
def skills_client(mock_genai_client):
17+
creds = mock.MagicMock()
18+
creds.token = "test_token"
19+
client = vertexai_client.Client(
20+
project="test-project", location="test-location", credentials=creds
21+
)
22+
client._genai_client = mock_genai_client
23+
return client.skills
24+
25+
26+
class TestGenaiSkills:
27+
mock_get_skill_response = {
28+
"name": "projects/test-project/locations/test-location/skills/test-skill",
29+
"displayName": "My Test Skill",
30+
# Add other expected fields from the Skill proto
31+
}
32+
33+
def test_get_skill(self, skills_client, mock_genai_client):
34+
"""Tests the get_skill method."""
35+
mock_genai_client.post.return_value = deps.Response(
36+
result=self.mock_get_skill_response,
37+
request=mock.MagicMock(),
38+
response=mock.MagicMock(),
39+
)
40+
41+
skill_name = "projects/test-project/locations/test-location/skills/test-skill"
42+
skill = skills_client.get(name=skill_name)
43+
44+
mock_genai_client.post.assert_called_once()
45+
call_args = mock_genai_client.post.call_args
46+
assert call_args[0][0] == skill_name
47+
assert call_args[1]["method"] == "GET"
48+
49+
assert isinstance(skill, genai.types.Skill)
50+
assert skill.name == skill_name
51+
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.BaseConfigOrDict] = 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.BaseConfigOrDict] = 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

0 commit comments

Comments
 (0)