Skip to content

Commit dd3df9b

Browse files
wanlin31copybara-github
authored andcommitted
feat: return headers for list method in all modules.
PiperOrigin-RevId: 785947079
1 parent 235ee99 commit dd3df9b

File tree

12 files changed

+128
-14
lines changed

12 files changed

+128
-14
lines changed

google/genai/batches.py

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1622,6 +1622,11 @@ def _ListBatchJobsResponse_from_mldev(
16221622
parent_object: Optional[dict[str, Any]] = None,
16231623
) -> dict[str, Any]:
16241624
to_object: dict[str, Any] = {}
1625+
if getv(from_object, ['sdkHttpResponse']) is not None:
1626+
setv(
1627+
to_object, ['sdk_http_response'], getv(from_object, ['sdkHttpResponse'])
1628+
)
1629+
16251630
if getv(from_object, ['nextPageToken']) is not None:
16261631
setv(to_object, ['next_page_token'], getv(from_object, ['nextPageToken']))
16271632

@@ -1784,6 +1789,11 @@ def _ListBatchJobsResponse_from_vertex(
17841789
parent_object: Optional[dict[str, Any]] = None,
17851790
) -> dict[str, Any]:
17861791
to_object: dict[str, Any] = {}
1792+
if getv(from_object, ['sdkHttpResponse']) is not None:
1793+
setv(
1794+
to_object, ['sdk_http_response'], getv(from_object, ['sdkHttpResponse'])
1795+
)
1796+
17871797
if getv(from_object, ['nextPageToken']) is not None:
17881798
setv(to_object, ['next_page_token'], getv(from_object, ['nextPageToken']))
17891799

@@ -2091,7 +2101,9 @@ def _list(
20912101
return_value = types.ListBatchJobsResponse._from_response(
20922102
response=response_dict, kwargs=parameter_model.model_dump()
20932103
)
2094-
2104+
return_value.sdk_http_response = types.HttpResponse(
2105+
headers=response.headers
2106+
)
20952107
self._api_client._verify_response(return_value)
20962108
return return_value
20972109

@@ -2522,7 +2534,9 @@ async def _list(
25222534
return_value = types.ListBatchJobsResponse._from_response(
25232535
response=response_dict, kwargs=parameter_model.model_dump()
25242536
)
2525-
2537+
return_value.sdk_http_response = types.HttpResponse(
2538+
headers=response.headers
2539+
)
25262540
self._api_client._verify_response(return_value)
25272541
return return_value
25282542

google/genai/caches.py

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1327,6 +1327,11 @@ def _ListCachedContentsResponse_from_mldev(
13271327
parent_object: Optional[dict[str, Any]] = None,
13281328
) -> dict[str, Any]:
13291329
to_object: dict[str, Any] = {}
1330+
if getv(from_object, ['sdkHttpResponse']) is not None:
1331+
setv(
1332+
to_object, ['sdk_http_response'], getv(from_object, ['sdkHttpResponse'])
1333+
)
1334+
13301335
if getv(from_object, ['nextPageToken']) is not None:
13311336
setv(to_object, ['next_page_token'], getv(from_object, ['nextPageToken']))
13321337

@@ -1386,6 +1391,11 @@ def _ListCachedContentsResponse_from_vertex(
13861391
parent_object: Optional[dict[str, Any]] = None,
13871392
) -> dict[str, Any]:
13881393
to_object: dict[str, Any] = {}
1394+
if getv(from_object, ['sdkHttpResponse']) is not None:
1395+
setv(
1396+
to_object, ['sdk_http_response'], getv(from_object, ['sdkHttpResponse'])
1397+
)
1398+
13891399
if getv(from_object, ['nextPageToken']) is not None:
13901400
setv(to_object, ['next_page_token'], getv(from_object, ['nextPageToken']))
13911401

@@ -1773,7 +1783,9 @@ def _list(
17731783
return_value = types.ListCachedContentsResponse._from_response(
17741784
response=response_dict, kwargs=parameter_model.model_dump()
17751785
)
1776-
1786+
return_value.sdk_http_response = types.HttpResponse(
1787+
headers=response.headers
1788+
)
17771789
self._api_client._verify_response(return_value)
17781790
return return_value
17791791

@@ -2165,7 +2177,9 @@ async def _list(
21652177
return_value = types.ListCachedContentsResponse._from_response(
21662178
response=response_dict, kwargs=parameter_model.model_dump()
21672179
)
2168-
2180+
return_value.sdk_http_response = types.HttpResponse(
2181+
headers=response.headers
2182+
)
21692183
self._api_client._verify_response(return_value)
21702184
return return_value
21712185

google/genai/files.py

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -267,6 +267,11 @@ def _ListFilesResponse_from_mldev(
267267
parent_object: Optional[dict[str, Any]] = None,
268268
) -> dict[str, Any]:
269269
to_object: dict[str, Any] = {}
270+
if getv(from_object, ['sdkHttpResponse']) is not None:
271+
setv(
272+
to_object, ['sdk_http_response'], getv(from_object, ['sdkHttpResponse'])
273+
)
274+
270275
if getv(from_object, ['nextPageToken']) is not None:
271276
setv(to_object, ['next_page_token'], getv(from_object, ['nextPageToken']))
272277

@@ -370,7 +375,9 @@ def _list(
370375
return_value = types.ListFilesResponse._from_response(
371376
response=response_dict, kwargs=parameter_model.model_dump()
372377
)
373-
378+
return_value.sdk_http_response = types.HttpResponse(
379+
headers=response.headers
380+
)
374381
self._api_client._verify_response(return_value)
375382
return return_value
376383

@@ -848,7 +855,9 @@ async def _list(
848855
return_value = types.ListFilesResponse._from_response(
849856
response=response_dict, kwargs=parameter_model.model_dump()
850857
)
851-
858+
return_value.sdk_http_response = types.HttpResponse(
859+
headers=response.headers
860+
)
852861
self._api_client._verify_response(return_value)
853862
return return_value
854863

google/genai/models.py

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3737,6 +3737,11 @@ def _ListModelsResponse_from_mldev(
37373737
parent_object: Optional[dict[str, Any]] = None,
37383738
) -> dict[str, Any]:
37393739
to_object: dict[str, Any] = {}
3740+
if getv(from_object, ['sdkHttpResponse']) is not None:
3741+
setv(
3742+
to_object, ['sdk_http_response'], getv(from_object, ['sdkHttpResponse'])
3743+
)
3744+
37403745
if getv(from_object, ['nextPageToken']) is not None:
37413746
setv(to_object, ['next_page_token'], getv(from_object, ['nextPageToken']))
37423747

@@ -4499,6 +4504,11 @@ def _ListModelsResponse_from_vertex(
44994504
parent_object: Optional[dict[str, Any]] = None,
45004505
) -> dict[str, Any]:
45014506
to_object: dict[str, Any] = {}
4507+
if getv(from_object, ['sdkHttpResponse']) is not None:
4508+
setv(
4509+
to_object, ['sdk_http_response'], getv(from_object, ['sdkHttpResponse'])
4510+
)
4511+
45024512
if getv(from_object, ['nextPageToken']) is not None:
45034513
setv(to_object, ['next_page_token'], getv(from_object, ['nextPageToken']))
45044514

@@ -5261,7 +5271,9 @@ def _list(
52615271
return_value = types.ListModelsResponse._from_response(
52625272
response=response_dict, kwargs=parameter_model.model_dump()
52635273
)
5264-
5274+
return_value.sdk_http_response = types.HttpResponse(
5275+
headers=response.headers
5276+
)
52655277
self._api_client._verify_response(return_value)
52665278
return return_value
52675279

@@ -6876,7 +6888,9 @@ async def _list(
68766888
return_value = types.ListModelsResponse._from_response(
68776889
response=response_dict, kwargs=parameter_model.model_dump()
68786890
)
6879-
6891+
return_value.sdk_http_response = types.HttpResponse(
6892+
headers=response.headers
6893+
)
68806894
self._api_client._verify_response(return_value)
68816895
return return_value
68826896

google/genai/pagers.py

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,8 @@
1818
# pylint: disable=protected-access
1919

2020
import copy
21-
from typing import Any, AsyncIterator,Awaitable, Callable, Generic, Iterator, Literal, TypeVar
21+
from typing import Any, AsyncIterator,Awaitable, Callable, Generic, Iterator, Literal, TypeVar, Union
22+
from . import types
2223

2324
T = TypeVar('T')
2425

@@ -43,6 +44,8 @@ def _init_page(
4344
self._page = getattr(response, self._name) or []
4445
self._idx = 0
4546

47+
self._sdk_http_response = getattr(response, 'sdk_http_response', None)
48+
4649
if not config:
4750
request_config = {}
4851
elif isinstance(config, dict):
@@ -110,6 +113,13 @@ def page_size(self) -> int:
110113

111114
return self._page_size
112115

116+
@property
117+
def sdk_http_response(self) -> Union[types.HttpResponse, None]:
118+
"""Returns the http response of the API response.
119+
"""
120+
121+
return self._sdk_http_response
122+
113123
@property
114124
def config(self) -> dict[str, Any]:
115125
"""Returns the configuration when making the API request for the next page.

google/genai/tests/batches/test_list.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@
5555

5656
def test_pager(client):
5757
batch_jobs = client.batches.list(config={'page_size': 10})
58+
assert 'content-type' in batch_jobs.sdk_http_response.headers
5859
assert batch_jobs.name == 'batch_jobs'
5960
assert batch_jobs.page_size == 10
6061
assert len(batch_jobs) <= 10
@@ -68,6 +69,7 @@ def test_pager(client):
6869
async def test_async_pager(client):
6970
batch_jobs = await client.aio.batches.list(config={'page_size': 10})
7071

72+
assert 'Content-Type' in batch_jobs.sdk_http_response.headers
7173
assert batch_jobs.name == 'batch_jobs'
7274
assert batch_jobs.page_size == 10
7375
assert len(batch_jobs) <= 10

google/genai/tests/caches/test_list.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@
4040

4141
def test_pager(client):
4242
cached_contents = client.caches.list(config={'page_size': 2})
43-
43+
assert 'content-type' in cached_contents.sdk_http_response.headers
4444
assert cached_contents.name == 'cached_contents'
4545
assert cached_contents.page_size == 2
4646
assert len(cached_contents) <= 2
@@ -56,6 +56,7 @@ def test_pager(client):
5656
async def test_async_pager(client):
5757
cached_contents = await client.aio.caches.list(config={'page_size': 2})
5858

59+
assert 'Content-Type' in cached_contents.sdk_http_response.headers
5960
assert cached_contents.name == 'cached_contents'
6061
assert cached_contents.page_size == 2
6162
assert len(cached_contents) <= 2

google/genai/tests/files/test_list.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@
4343
def test_pager(client):
4444
with pytest_helper.exception_if_vertex(client, ValueError):
4545
files = client.files.list(config={'page_size': 2})
46-
46+
assert 'content-type' in files.sdk_http_response.headers
4747
assert files.name == 'files'
4848
assert files.page_size == 2
4949
assert len(files) <= 2
@@ -60,6 +60,7 @@ async def test_async_pager(client):
6060
with pytest_helper.exception_if_vertex(client, ValueError):
6161
files = await client.aio.files.list(config={'page_size': 2})
6262

63+
assert 'Content-Type' in files.sdk_http_response.headers
6364
assert files.name == 'files'
6465
assert files.page_size == 2
6566
assert len(files) <= 2

google/genai/tests/models/test_list.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,8 @@ def test_base_models_pager(client):
9292

9393
# Iterate through all the pages. Then next_page() should raise an exception.
9494
for _ in pager:
95+
assert 'content-type' in pager.sdk_http_response.headers
96+
assert 'content-encoding' in pager.sdk_http_response.headers
9597
pass
9698
with pytest.raises(IndexError, match='No more pages to fetch.'):
9799
pager.next_page()
@@ -193,6 +195,8 @@ def test_empty_api_response_empty_dict_headers(mock_api_client, client):
193195
async def test_tuned_models_async_pager(client):
194196
pager = await client.aio.models.list(config={'page_size': 10, 'query_base': False})
195197

198+
assert 'Content-Type' in pager.sdk_http_response.headers
199+
assert 'Content-Encoding' in pager.sdk_http_response.headers
196200
assert pager.name == 'models'
197201
assert pager.page_size == 10
198202
assert len(pager) <= 10

google/genai/tests/tunings/test_list.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@
4747

4848
def test_pager(client):
4949
tuning_jobs = client.tunings.list(config={'page_size': 2})
50-
50+
assert 'content-type' in tuning_jobs.sdk_http_response.headers
5151
assert tuning_jobs.name == 'tuning_jobs'
5252
assert tuning_jobs.page_size == 2
5353
assert len(tuning_jobs) <= 2
@@ -63,6 +63,7 @@ def test_pager(client):
6363
async def test_async_pager(client):
6464
tuning_jobs = await client.aio.tunings.list(config={'page_size': 2})
6565

66+
assert 'Content-Type' in tuning_jobs.sdk_http_response.headers
6667
assert tuning_jobs.name == 'tuning_jobs'
6768
assert tuning_jobs.page_size == 2
6869
assert len(tuning_jobs) <= 2

0 commit comments

Comments
 (0)