diff --git a/agentplatform/_genai/skill_revisions.py b/agentplatform/_genai/skill_revisions.py index 11484e428b..3f45554faf 100644 --- a/agentplatform/_genai/skill_revisions.py +++ b/agentplatform/_genai/skill_revisions.py @@ -15,6 +15,7 @@ # Code generated by the Google Gen AI SDK generator DO NOT EDIT. +from functools import partial import json import logging from typing import Any, Optional, Union @@ -24,6 +25,7 @@ from google.genai import _common from google.genai._common import get_value_by_path as getv from google.genai._common import set_value_by_path as setv +from google.genai.pagers import AsyncPager, Pager from . import types @@ -151,16 +153,12 @@ def get( self._api_client._verify_response(return_value) return return_value - def list( + def _list( self, *, name: str, config: Optional[types.ListSkillRevisionsConfigOrDict] = None, ) -> types.ListSkillRevisionsResponse: - """ - Lists Skill Revisions. - """ - parameter_model = types._ListSkillRevisionsRequestParameters( name=name, config=config, @@ -225,6 +223,24 @@ def list( self._api_client._verify_response(return_value) return return_value + def list( + self, + *, + name: str, + config: Optional[types.ListSkillRevisionsConfigOrDict] = None, + ) -> Pager[types.SkillRevision]: + """ + Lists Skill Revisions. + """ + + list_request = partial(self._list, name=name) + return Pager( + "skill_revisions", + list_request, + self._list(name=name, config=config), + config, + ) + class AsyncSkillRevisions(_api_module.BaseModule): """Class for managing Skill Revisions in the Skill Registry.""" @@ -300,16 +316,12 @@ async def get( self._api_client._verify_response(return_value) return return_value - async def list( + async def _list( self, *, name: str, config: Optional[types.ListSkillRevisionsConfigOrDict] = None, ) -> types.ListSkillRevisionsResponse: - """ - Lists Skill Revisions. - """ - parameter_model = types._ListSkillRevisionsRequestParameters( name=name, config=config, @@ -375,3 +387,17 @@ async def list( self._api_client._verify_response(return_value) return return_value + + async def list( + self, + *, + name: str, + config: Optional[types.ListSkillRevisionsConfigOrDict] = None, + ) -> AsyncPager[types.SkillRevision]: + list_request = partial(self._list, name=name) + return AsyncPager( + "skill_revisions", + list_request, + await self._list(name=name, config=config), + config, + ) diff --git a/agentplatform/_genai/types/common.py b/agentplatform/_genai/types/common.py index 55e841d973..b664633f96 100644 --- a/agentplatform/_genai/types/common.py +++ b/agentplatform/_genai/types/common.py @@ -18468,6 +18468,10 @@ class ListSkillsConfig(_common.BaseModel): filter: Optional[str] = Field( default=None, description="""Optional. The standard list filter.""" ) + max_results: Optional[int] = Field( + default=None, + description="""Optional. Maximum number of results to return. Defaults to 100.""", + ) class ListSkillsConfigDict(TypedDict, total=False): @@ -18485,6 +18489,9 @@ class ListSkillsConfigDict(TypedDict, total=False): filter: Optional[str] """Optional. The standard list filter.""" + max_results: Optional[int] + """Optional. Maximum number of results to return. Defaults to 100.""" + ListSkillsConfigOrDict = Union[ListSkillsConfig, ListSkillsConfigDict] @@ -18752,6 +18759,10 @@ class ListSkillRevisionsConfig(_common.BaseModel): ) page_size: Optional[int] = Field(default=None, description="""""") page_token: Optional[str] = Field(default=None, description="""""") + max_results: Optional[int] = Field( + default=None, + description="""Optional. Maximum number of results to return. Defaults to 100.""", + ) class ListSkillRevisionsConfigDict(TypedDict, total=False): @@ -18765,6 +18776,9 @@ class ListSkillRevisionsConfigDict(TypedDict, total=False): page_token: Optional[str] """""" + max_results: Optional[int] + """Optional. Maximum number of results to return. Defaults to 100.""" + ListSkillRevisionsConfigOrDict = Union[ ListSkillRevisionsConfig, ListSkillRevisionsConfigDict diff --git a/tests/unit/agentplatform/genai/replays/test_skills_revisions_list.py b/tests/unit/agentplatform/genai/replays/test_skills_revisions_list.py index 99cf9bd19f..406fe0656d 100644 --- a/tests/unit/agentplatform/genai/replays/test_skills_revisions_list.py +++ b/tests/unit/agentplatform/genai/replays/test_skills_revisions_list.py @@ -40,8 +40,8 @@ def test_list_skill_revisions(client, tmp_path): assert created_skill.name is not None # 2. List revisions - revisions_response = client.skills.revisions.list(name=created_skill.name) - revisions_list = revisions_response.skill_revisions + pager = client.skills.revisions.list(name=created_skill.name) + revisions_list = list(pager) assert len(revisions_list) > 0 first_revision = revisions_list[0] diff --git a/tests/unit/agentplatform/genai/test_genai_skills.py b/tests/unit/agentplatform/genai/test_genai_skills.py index ee3fb71c93..d640125c56 100644 --- a/tests/unit/agentplatform/genai/test_genai_skills.py +++ b/tests/unit/agentplatform/genai/test_genai_skills.py @@ -969,11 +969,15 @@ def test_list_skill_revisions(self, skills_client): body=json.dumps(mock_response) ) - response = skills_client.revisions.list(name=skill_name) + pager = skills_client.revisions.list(name=skill_name) - assert isinstance(response, genai.types.ListSkillRevisionsResponse) - assert len(response.skill_revisions) == 1 - assert response.skill_revisions[0].name == f"{skill_name}/revisions/rev-1" + from google.genai import pagers + + assert isinstance(pager, pagers.Pager) + + revisions = list(pager) + assert len(revisions) == 1 + assert revisions[0].name == f"{skill_name}/revisions/rev-1" request_mock.assert_called_once_with( "get", @@ -1001,11 +1005,18 @@ async def test_list_skill_revisions_async(self, async_skills_client): body=json.dumps(mock_response) ) - response = await async_skills_client.revisions.list(name=skill_name) + pager = await async_skills_client.revisions.list(name=skill_name) + + from google.genai import pagers + + assert isinstance(pager, pagers.AsyncPager) - assert isinstance(response, genai.types.ListSkillRevisionsResponse) - assert len(response.skill_revisions) == 1 - assert response.skill_revisions[0].name == f"{skill_name}/revisions/rev-1" + revisions = [] + async for revision in pager: + revisions.append(revision) + + assert len(revisions) == 1 + assert revisions[0].name == f"{skill_name}/revisions/rev-1" request_mock.assert_called_once_with( "get", @@ -1013,3 +1024,340 @@ async def test_list_skill_revisions_async(self, async_skills_client): {"_url": {"name": skill_name}}, None, ) + + def test_list_skills_with_max_results(self, skills_client): + mock_list_response = { + "skills": [ + { + "name": "projects/test-project/locations/test-location/skills/skill-1", + "displayName": "Skill 1", + }, + { + "name": "projects/test-project/locations/test-location/skills/skill-2", + "displayName": "Skill 2", + }, + ], + "nextPageToken": "token-123", + } + mock_list_response_page_2 = { + "skills": [ + { + "name": "projects/test-project/locations/test-location/skills/skill-3", + "displayName": "Skill 3", + }, + { + "name": "projects/test-project/locations/test-location/skills/skill-4", + "displayName": "Skill 4", + }, + ], + "nextPageToken": "token-456", + } + + with mock.patch.object( + skills_client._api_client, "request", autospec=True + ) as request_mock: + request_mock.side_effect = [ + genai_types.HttpResponse(body=json.dumps(mock_list_response)), + genai_types.HttpResponse(body=json.dumps(mock_list_response_page_2)), + ] + + pager = skills_client.list(config={"max_results": 3}) + skills = list(pager) + + assert len(skills) == 3 + assert skills[0].display_name == "Skill 1" + assert skills[1].display_name == "Skill 2" + assert skills[2].display_name == "Skill 3" + + assert request_mock.call_count == 2 + request_mock.assert_has_calls( + [ + mock.call("get", "skills", {}, None), + mock.call( + "get", + "skills?pageToken=token-123", + {"_query": {"pageToken": "token-123"}}, + None, + ), + ] + ) + + @pytest.mark.asyncio + async def test_list_skills_with_max_results_async(self, async_skills_client): + mock_list_response = { + "skills": [ + { + "name": "projects/test-project/locations/test-location/skills/skill-1", + "displayName": "Skill 1", + }, + { + "name": "projects/test-project/locations/test-location/skills/skill-2", + "displayName": "Skill 2", + }, + ], + "nextPageToken": "token-123", + } + mock_list_response_page_2 = { + "skills": [ + { + "name": "projects/test-project/locations/test-location/skills/skill-3", + "displayName": "Skill 3", + }, + { + "name": "projects/test-project/locations/test-location/skills/skill-4", + "displayName": "Skill 4", + }, + ], + "nextPageToken": "token-456", + } + + with mock.patch.object( + async_skills_client._api_client, "async_request", autospec=True + ) as request_mock: + request_mock.side_effect = [ + genai_types.HttpResponse(body=json.dumps(mock_list_response)), + genai_types.HttpResponse(body=json.dumps(mock_list_response_page_2)), + ] + + pager = await async_skills_client.list(config={"max_results": 3}) + skills = [] + async for skill in pager: + skills.append(skill) + + assert len(skills) == 3 + assert skills[0].display_name == "Skill 1" + assert skills[1].display_name == "Skill 2" + assert skills[2].display_name == "Skill 3" + + assert request_mock.call_count == 2 + request_mock.assert_has_calls( + [ + mock.call("get", "skills", {}, None), + mock.call( + "get", + "skills?pageToken=token-123", + {"_query": {"pageToken": "token-123"}}, + None, + ), + ] + ) + + def test_list_skill_revisions_with_max_results(self, skills_client): + skill_name = "projects/test-project/locations/test-location/skills/test-skill" + mock_response = { + "skillRevisions": [ + {"name": f"{skill_name}/revisions/rev-1", "state": "ACTIVE"}, + {"name": f"{skill_name}/revisions/rev-2", "state": "ACTIVE"}, + ], + "nextPageToken": "token-123", + } + mock_response_page_2 = { + "skillRevisions": [ + {"name": f"{skill_name}/revisions/rev-3", "state": "ACTIVE"}, + ], + } + + with mock.patch.object( + skills_client._api_client, "request", autospec=True + ) as request_mock: + request_mock.side_effect = [ + genai_types.HttpResponse(body=json.dumps(mock_response)), + genai_types.HttpResponse(body=json.dumps(mock_response_page_2)), + ] + + pager = skills_client.revisions.list( + name=skill_name, config={"max_results": 2} + ) + revisions = list(pager) + + assert len(revisions) == 2 + assert revisions[0].name == f"{skill_name}/revisions/rev-1" + assert revisions[1].name == f"{skill_name}/revisions/rev-2" + + assert request_mock.call_count == 1 + request_mock.assert_called_once_with( + "get", + f"{skill_name}/revisions", + {"_url": {"name": skill_name}}, + None, + ) + + @pytest.mark.asyncio + async def test_list_skill_revisions_with_max_results_async( + self, async_skills_client + ): + skill_name = "projects/test-project/locations/test-location/skills/test-skill" + mock_response = { + "skillRevisions": [ + {"name": f"{skill_name}/revisions/rev-1", "state": "ACTIVE"}, + {"name": f"{skill_name}/revisions/rev-2", "state": "ACTIVE"}, + ], + "nextPageToken": "token-123", + } + mock_response_page_2 = { + "skillRevisions": [ + {"name": f"{skill_name}/revisions/rev-3", "state": "ACTIVE"}, + ], + } + + with mock.patch.object( + async_skills_client._api_client, "async_request", autospec=True + ) as request_mock: + request_mock.side_effect = [ + genai_types.HttpResponse(body=json.dumps(mock_response)), + genai_types.HttpResponse(body=json.dumps(mock_response_page_2)), + ] + + pager = await async_skills_client.revisions.list( + name=skill_name, config={"max_results": 2} + ) + revisions = [] + async for revision in pager: + revisions.append(revision) + + assert len(revisions) == 2 + assert revisions[0].name == f"{skill_name}/revisions/rev-1" + assert revisions[1].name == f"{skill_name}/revisions/rev-2" + + assert request_mock.call_count == 1 + request_mock.assert_called_once_with( + "get", + f"{skill_name}/revisions", + {"_url": {"name": skill_name}}, + None, + ) + + def test_pager_invalid_max_results(self, skills_client): + import pydantic + + mock_list_response = {"skills": []} + with mock.patch.object( + skills_client._api_client, "request", autospec=True + ) as request_mock: + request_mock.return_value = genai_types.HttpResponse( + body=json.dumps(mock_list_response) + ) + + with pytest.raises(ValueError, match="max_results must be greater than 0"): + list(skills_client.list(config={"max_results": 0})) + + with pytest.raises(ValueError, match="max_results must be greater than 0"): + list(skills_client.list(config={"max_results": -1})) + + with pytest.raises(pydantic.ValidationError): + list(skills_client.list(config={"max_results": "invalid"})) + + def test_list_skills_default_limit(self, skills_client): + # 50 pages of size 2 = 100 items (reaches default limit) + responses = [] + for i in range(50): + responses.append(genai_types.HttpResponse(body=json.dumps({ + "skills": [ + {"name": f"projects/test-project/locations/test-location/skills/skill-{i*2+1}", "displayName": f"Skill {i*2+1}"}, + {"name": f"projects/test-project/locations/test-location/skills/skill-{i*2+2}", "displayName": f"Skill {i*2+2}"}, + ], + "nextPageToken": f"token-{i+1}" + }))) + # 51st page (should NOT be fetched) + responses.append(genai_types.HttpResponse(body=json.dumps({ + "skills": [ + {"name": "projects/test-project/locations/test-location/skills/skill-101", "displayName": "Skill 101"}, + ] + }))) + + with mock.patch.object( + skills_client._api_client, "request", autospec=True + ) as request_mock: + request_mock.side_effect = responses + + # No max_results specified, should default to 100 + pager = skills_client.list() + skills = list(pager) + + assert len(skills) == 100 + assert skills[0].display_name == "Skill 1" + assert skills[99].display_name == "Skill 100" + + # Should have called request 50 times (fetching 100 items) + assert request_mock.call_count == 50 + # Ensure the last call was with the 49th token + request_mock.assert_has_calls( + [ + mock.call("get", "skills", {}, None), + ] + [ + mock.call( + "get", + f"skills?pageToken=token-{i}", + {"_query": {"pageToken": f"token-{i}"}}, + None, + ) for i in range(1, 50) + ] + ) + + def test_list_skill_revisions_default_limit(self, skills_client): + skill_name = "projects/test-project/locations/test-location/skills/test-skill" + # 50 pages of size 2 = 100 items + responses = [] + for i in range(50): + responses.append(genai_types.HttpResponse(body=json.dumps({ + "skillRevisions": [ + {"name": f"{skill_name}/revisions/rev-{i*2+1}", "state": "ACTIVE"}, + {"name": f"{skill_name}/revisions/rev-{i*2+2}", "state": "ACTIVE"}, + ], + "nextPageToken": f"token-{i+1}" + }))) + # 51st page (should NOT be fetched) + responses.append(genai_types.HttpResponse(body=json.dumps({ + "skillRevisions": [ + {"name": f"{skill_name}/revisions/rev-101", "state": "ACTIVE"}, + ] + }))) + + with mock.patch.object( + skills_client._api_client, "request", autospec=True + ) as request_mock: + request_mock.side_effect = responses + + pager = skills_client.revisions.list(name=skill_name) + revisions = list(pager) + + assert len(revisions) == 100 + assert revisions[0].name == f"{skill_name}/revisions/rev-1" + assert revisions[99].name == f"{skill_name}/revisions/rev-100" + + assert request_mock.call_count == 50 + + @pytest.mark.asyncio + async def test_list_skills_default_limit_async(self, async_skills_client): + # 50 pages of size 2 = 100 items + responses = [] + for i in range(50): + responses.append(genai_types.HttpResponse(body=json.dumps({ + "skills": [ + {"name": f"projects/test-project/locations/test-location/skills/skill-{i*2+1}", "displayName": f"Skill {i*2+1}"}, + {"name": f"projects/test-project/locations/test-location/skills/skill-{i*2+2}", "displayName": f"Skill {i*2+2}"}, + ], + "nextPageToken": f"token-{i+1}" + }))) + # 51st page (should NOT be fetched) + responses.append(genai_types.HttpResponse(body=json.dumps({ + "skills": [ + {"name": "projects/test-project/locations/test-location/skills/skill-101", "displayName": "Skill 101"}, + ] + }))) + + with mock.patch.object( + async_skills_client._api_client, "async_request", autospec=True + ) as request_mock: + request_mock.side_effect = responses + + pager = await async_skills_client.list() + skills = [] + async for skill in pager: + skills.append(skill) + + assert len(skills) == 100 + assert skills[0].display_name == "Skill 1" + assert skills[99].display_name == "Skill 100" + + assert request_mock.call_count == 50