diff --git a/google/genai/_api_client.py b/google/genai/_api_client.py index 3c4cbc247..27adcf813 100644 --- a/google/genai/_api_client.py +++ b/google/genai/_api_client.py @@ -34,7 +34,7 @@ import sys import threading import time -from typing import Any, AsyncIterator, Iterator, Optional, TYPE_CHECKING, Tuple, Union +from typing import Any, AsyncIterator, Iterator, Optional, TYPE_CHECKING, Tuple, Union, cast, overload from urllib.parse import urlparse from urllib.parse import urlunparse import warnings @@ -638,6 +638,9 @@ def __del__(self) -> None: class BaseApiClient: """Client for calling HTTP APIs sending and receiving JSON.""" + vertexai: Optional[bool] = None + custom_base_url: Optional[str] = None + def __init__( self, vertexai: Optional[bool] = None, @@ -1962,21 +1965,55 @@ def _upload_fd( raise ValueError('Failed to upload file: Upload status is not finalized.') return HttpResponse(response.headers, response_stream=[response.text]) + @overload def download_file( self, path: str, *, http_options: Optional[HttpOptionsOrDict] = None, - ) -> Union[Any, bytes]: + destination: None = None, + chunk_size: int = 1024 * 1024, + ) -> bytes: + ... + + @overload + def download_file( + self, + path: str, + *, + http_options: Optional[HttpOptionsOrDict] = None, + destination: Union[str, os.PathLike[str], io.IOBase], + chunk_size: int = 1024 * 1024, + ) -> None: + ... + + def download_file( + self, + path: str, + *, + http_options: Optional[HttpOptionsOrDict] = None, + destination: Optional[Union[str, os.PathLike[str], io.IOBase]] = None, + chunk_size: int = 1024 * 1024, + ) -> Optional[bytes]: """Downloads the file data. Args: path: The request path with query params. http_options: The http options to use for the request. + destination: Optional local file path or writable stream. + chunk_size: The chunk size in bytes to stream. - returns: - The file bytes + Returns: + The file bytes if destination is None, otherwise None. """ + if destination is not None and not isinstance( + destination, (str, os.PathLike) + ) and not hasattr(destination, 'write'): + raise ValueError( + f'Unsupported destination type: {type(destination)}. ' + 'Expected str, os.PathLike, or a writable file-like object.' + ) + http_request = self._build_request( 'get', path=path, request_dict={}, http_options=http_options ) @@ -1988,18 +2025,86 @@ def download_file( else: data = http_request.data - response = self._httpx_client.request( # type: ignore[union-attr] - method=http_request.method, - url=http_request.url, - headers=http_request.headers, - content=data, - timeout=http_request.timeout, - ) + def _write_chunks(chunks: Iterator[bytes]) -> None: + if isinstance(destination, (str, os.PathLike)): + with open(destination, 'wb') as f: + for chunk in chunks: + f.write(chunk) + elif destination is not None and hasattr(destination, 'write'): + for chunk in chunks: + destination.write(chunk) - errors.APIError.raise_for_response(response) - return HttpResponse( - response.headers, byte_stream=[response.read()] - ).byte_stream[0] + if self._use_google_auth_sync(): + url = str(http_request.url) + if self._authorized_session is None: + from google.auth.transport.requests import AuthorizedSession # pylint: disable=g-import-not-at-top + + self._authorized_session = AuthorizedSession( # type: ignore[no-untyped-call] + self._credentials, + max_refresh_attempts=1, + ) + client_cert_source = mtls.default_client_cert_source() # type: ignore[no-untyped-call] + self._authorized_session.configure_mtls_channel( + client_cert_source + ) # type: ignore[no-untyped-call] + if self._authorized_session._is_mtls and 'googleapis.com' in url: + if 'sandbox' in url: + url = url.replace( + 'sandbox.googleapis.com', 'mtls.sandbox.googleapis.com' + ) + else: + url = url.replace('googleapis.com', 'mtls.googleapis.com') + if destination is not None: + response = self._authorized_session.request( # type: ignore[no-untyped-call] + method=http_request.method.upper(), + url=url, + data=data, + headers=http_request.headers, + timeout=http_request.timeout, + stream=True, + ) + try: + errors.APIError.raise_for_response(response) + _write_chunks(response.iter_content(chunk_size=chunk_size)) + finally: + response.close() + return None + else: + response = self._authorized_session.request( # type: ignore[no-untyped-call] + method=http_request.method.upper(), + url=url, + data=data, + headers=http_request.headers, + timeout=http_request.timeout, + ) + errors.APIError.raise_for_response(response) + return cast(bytes, response.content) + else: + if destination is not None: + httpx_request = self._httpx_client.build_request( # type: ignore[union-attr] + method=http_request.method, + url=http_request.url, + content=data, + headers=http_request.headers, + timeout=http_request.timeout, + ) + response = self._httpx_client.send(httpx_request, stream=True) # type: ignore[union-attr, arg-type] + try: + errors.APIError.raise_for_response(response) + _write_chunks(response.iter_bytes(chunk_size=chunk_size)) + finally: + response.close() + return None + else: + response = self._httpx_client.request( # type: ignore[union-attr] + method=http_request.method, + url=http_request.url, + content=data, + headers=http_request.headers, + timeout=http_request.timeout, + ) + errors.APIError.raise_for_response(response) + return cast(bytes, response.read()) async def async_upload_file( self, @@ -2237,21 +2342,55 @@ async def _async_upload_fd( client_response.headers, response_stream=[client_response.text] ) + @overload async def async_download_file( self, path: str, *, http_options: Optional[HttpOptionsOrDict] = None, - ) -> Union[Any, bytes]: - """Downloads the file data. + destination: None = None, + chunk_size: int = 1024 * 1024, + ) -> bytes: + ... + + @overload + async def async_download_file( + self, + path: str, + *, + http_options: Optional[HttpOptionsOrDict] = None, + destination: Union[str, os.PathLike[str], io.IOBase], + chunk_size: int = 1024 * 1024, + ) -> None: + ... + + async def async_download_file( + self, + path: str, + *, + http_options: Optional[HttpOptionsOrDict] = None, + destination: Optional[Union[str, os.PathLike[str], io.IOBase]] = None, + chunk_size: int = 1024 * 1024, + ) -> Optional[bytes]: + """Downloads the file data asynchronously. Args: path: The request path with query params. http_options: The http options to use for the request. + destination: Optional local file path or writable stream. + chunk_size: The chunk size in bytes to stream. - returns: - The file bytes + Returns: + The file bytes if destination is None, otherwise None. """ + if destination is not None and not isinstance( + destination, (str, os.PathLike) + ) and not hasattr(destination, 'write'): + raise ValueError( + f'Unsupported destination type: {type(destination)}. ' + 'Expected str, os.PathLike, or a writable file-like object.' + ) + http_request = self._build_request( 'get', path=path, request_dict={}, http_options=http_options ) @@ -2263,34 +2402,89 @@ async def async_download_file( else: data = http_request.data + async def _write_chunks(chunks: AsyncIterator[bytes]) -> None: + if isinstance(destination, (str, os.PathLike)): + with open(destination, 'wb') as f: + async for chunk in chunks: + f.write(chunk) + elif destination is not None and hasattr(destination, 'write'): + async for chunk in chunks: + res = destination.write(chunk) + if inspect.isawaitable(res): + await res + if self._use_aiohttp(): session = await self._get_aiohttp_session() # type: ignore[assignment] + url = http_request.url + if self._use_google_auth_async(): + client_cert_source = mtls.default_client_cert_source() # type: ignore[no-untyped-call] + await session.configure_mtls_channel( # type: ignore[union-attr] + client_cert_source + ) + if session._is_mtls and 'googleapis.com' in url: # type: ignore[union-attr] + if 'sandbox' in url: + url = url.replace( + 'sandbox.googleapis.com', 'mtls.sandbox.googleapis.com' + ) + else: + url = url.replace('googleapis.com', 'mtls.googleapis.com') response = await session.request( # type: ignore[union-attr] method=http_request.method, - url=http_request.url, + url=url, headers=http_request.headers, data=data, timeout=aiohttp.ClientTimeout(total=http_request.timeout), + **self._async_client_session_request_args, ) - await errors.APIError.raise_for_async_response(response) - - return HttpResponse( - response.headers, byte_stream=[await response.read()] - ).byte_stream[0] + if destination is not None: + try: + await errors.APIError.raise_for_async_response(response) + if hasattr(response, '_response'): + raw_response = response._response + else: + raw_response = response + await _write_chunks(raw_response.content.iter_chunked(chunk_size)) + finally: + response.close() + return None + else: + try: + await errors.APIError.raise_for_async_response(response) + return cast(bytes, await response.read()) + finally: + response.close() else: # aiohttp is not available. Fall back to httpx. - client_response = await self._async_httpx_client.request( # type: ignore[union-attr] - method=http_request.method, - url=http_request.url, - headers=http_request.headers, - content=data, - timeout=http_request.timeout, - ) - await errors.APIError.raise_for_async_response(client_response) - - return HttpResponse( - client_response.headers, byte_stream=[client_response.read()] - ).byte_stream[0] + if destination is not None: + httpx_request = self._async_httpx_client.build_request( # type: ignore[union-attr] + method=http_request.method, + url=http_request.url, + content=data, + headers=http_request.headers, + timeout=http_request.timeout, + ) + client_response = await self._async_httpx_client.send( # type: ignore[union-attr] + httpx_request, # type: ignore[arg-type] + stream=True, + ) + try: + await errors.APIError.raise_for_async_response(client_response) + await _write_chunks( + client_response.aiter_bytes(chunk_size=chunk_size) + ) + finally: + await client_response.aclose() + return None + else: + client_response = await self._async_httpx_client.request( # type: ignore[union-attr] + method=http_request.method, + url=http_request.url, + headers=http_request.headers, + content=data, + timeout=http_request.timeout, + ) + await errors.APIError.raise_for_async_response(client_response) + return cast(bytes, client_response.read()) # This method does nothing in the real api client. It is used in the # replay_api_client to verify the response from the SDK method matches the diff --git a/google/genai/_replay_api_client.py b/google/genai/_replay_api_client.py index efb29151d..90b0a8078 100644 --- a/google/genai/_replay_api_client.py +++ b/google/genai/_replay_api_client.py @@ -24,7 +24,7 @@ import json import os import re -from typing import Any, Literal, Optional, Union, Iterator, AsyncIterator +from typing import Any, AsyncIterator, Iterator, Literal, Optional, Union, cast, overload import google.auth @@ -410,7 +410,7 @@ def _record_interaction( headers=dict(http_response.headers), body_segments=list(http_response.segments()), byte_segments=[ - seg[:100] + b'...' for seg in http_response.byte_segments() + seg[:100] for seg in http_response.byte_segments() ], status_code=http_response.status_code, sdk_response_segments=[], @@ -426,7 +426,7 @@ def _record_interaction( response = ReplayResponse( headers={}, body_segments=[], - byte_segments=[http_response], + byte_segments=[http_response[:64 * 1024]], sdk_response_segments=[], ) else: @@ -672,34 +672,128 @@ async def async_upload_file( else: return self._build_response_from_replay(request) + @overload def download_file( - self, path: str, *, http_options: Optional[HttpOptionsOrDict] = None - ) -> Union[HttpResponse, bytes, Any]: + self, + path: str, + *, + http_options: Optional[HttpOptionsOrDict] = None, + destination: None = None, + chunk_size: int = 1024 * 1024, + ) -> bytes: + ... + + @overload + def download_file( + self, + path: str, + *, + http_options: Optional[HttpOptionsOrDict] = None, + destination: Union[str, os.PathLike[str], io.IOBase], + chunk_size: int = 1024 * 1024, + ) -> None: + ... + + def download_file( + self, + path: str, + *, + http_options: Optional[HttpOptionsOrDict] = None, + destination: Optional[Union[str, os.PathLike[str], io.IOBase]] = None, + chunk_size: int = 1024 * 1024, + ) -> Optional[bytes]: self._initialize_replay_session_if_not_loaded() request = self._build_request( 'get', path=path, request_dict={}, http_options=http_options ) if self._should_call_api(): with _record_on_api_error(self, request): - result = super().download_file(path, http_options=http_options) - self._record_interaction(request, result) - return result + content = super().download_file( + path, + http_options=http_options, + destination=None, + chunk_size=chunk_size, + ) + self._record_interaction(request, content) else: - return self._build_response_from_replay(request).byte_stream[0] + content = cast( + bytes, self._build_response_from_replay(request).byte_stream[0] + ) + + if destination is not None: + if isinstance(destination, (str, os.PathLike)): + with open(destination, 'wb') as f: + f.write(content) + elif hasattr(destination, 'write'): + destination.write(content) + else: + raise ValueError( + f'Unsupported destination type: {type(destination)}. ' + 'Expected str, os.PathLike, or a writable file-like object.' + ) + return None + return content + @overload async def async_download_file( - self, path: str, *, http_options: Optional[HttpOptionsOrDict] = None - ) -> Any: + self, + path: str, + *, + http_options: Optional[HttpOptionsOrDict] = None, + destination: None = None, + chunk_size: int = 1024 * 1024, + ) -> bytes: + ... + + @overload + async def async_download_file( + self, + path: str, + *, + http_options: Optional[HttpOptionsOrDict] = None, + destination: Union[str, os.PathLike[str], io.IOBase], + chunk_size: int = 1024 * 1024, + ) -> None: + ... + + async def async_download_file( + self, + path: str, + *, + http_options: Optional[HttpOptionsOrDict] = None, + destination: Optional[Union[str, os.PathLike[str], io.IOBase]] = None, + chunk_size: int = 1024 * 1024, + ) -> Optional[bytes]: self._initialize_replay_session_if_not_loaded() request = self._build_request( 'get', path=path, request_dict={}, http_options=http_options ) if self._should_call_api(): async with _async_record_on_api_error(self, request): - result = await super().async_download_file( - path, http_options=http_options + content = await super().async_download_file( + path, + http_options=http_options, + destination=None, + chunk_size=chunk_size, ) - self._record_interaction(request, result) - return result + self._record_interaction(request, content) else: - return self._build_response_from_replay(request).byte_stream[0] + content = cast( + bytes, self._build_response_from_replay(request).byte_stream[0] + ) + + if destination is not None: + if isinstance(destination, (str, os.PathLike)): + with open(destination, 'wb') as f: + f.write(content) + elif hasattr(destination, 'write'): + res = destination.write(content) + if inspect.isawaitable(res): + await res + else: + raise ValueError( + f'Unsupported destination type: {type(destination)}. ' + 'Expected str, os.PathLike, or a writable file-like object.' + ) + return None + return content diff --git a/google/genai/files.py b/google/genai/files.py index 271c066ca..e49de2999 100644 --- a/google/genai/files.py +++ b/google/genai/files.py @@ -20,7 +20,7 @@ import json import logging import os -from typing import Any, Optional, Union +from typing import Any, Optional, Union, overload from urllib.parse import urlencode import google.auth @@ -652,43 +652,74 @@ def upload( kwargs=config_model.model_dump() if config else {}, ) + # from typing import overload + + @overload def download( self, *, file: Union[str, types.File, types.Video, types.GeneratedVideo], + destination: None = None, config: Optional[types.DownloadFileConfigOrDict] = None, ) -> bytes: + ... + + @overload + def download( + self, + *, + file: Union[str, types.File, types.Video, types.GeneratedVideo], + destination: Union[str, os.PathLike[str], io.IOBase], + config: Optional[types.DownloadFileConfigOrDict] = None, + ) -> None: + ... + + def download( + self, + *, + file: Union[str, types.File, types.Video, types.GeneratedVideo], + destination: Optional[Union[str, os.PathLike[str], io.IOBase]] = None, + config: Optional[types.DownloadFileConfigOrDict] = None, + ) -> Optional[bytes]: """Downloads a file's data from storage. Files created by `upload` can't be downloaded. You can tell which files are downloadable by checking the `source` or `download_uri` property. - Note: This method returns the data as bytes. For `Video` and - `GeneratedVideo` objects there is an additional side effect, that it also - sets the `video_bytes` property on the `Video` object. + Note: When destination is None, this method returns the data as bytes. For + `Video` and `GeneratedVideo` objects, if destination is None, it also sets + the + `video_bytes` property on the `Video` object. When destination is provided, + chunks are streamed directly to disk or the writable stream without loading + the full file into memory, and None is returned. Args: - file (str): A file name, uri, or file object. Identifying which file to - download. - config (DownloadFileConfigOrDict): Optional, configuration for the get - method. + file: A file name, uri, or file object identifying which file to download. + destination: Optional, a local file path or writable binary stream to + write the file directly to in chunks, avoiding holding the full file in + memory. + config: Optional, configuration for the get method. Returns: - File: The file data as bytes. + bytes: The file data as bytes if destination is None; otherwise None. Usage: .. code-block:: python - for file client.files.list(): + for file in client.files.list(): if file.download_uri is not None: break else: raise ValueError('No files found with a `download_uri`.') + + # Download to memory: data = client.files.download(file=file) - # data = client.files.download(file=file.name) - # data = client.files.download(file=file.download_uri) + # Stream directly to a file path (avoids memory spike): + client.files.download(file=file, destination="output.mp4") + + # Download video object: video = types.Video(uri=file.uri) video_bytes = client.files.download(file=video) video.video_bytes @@ -724,12 +755,14 @@ def download( data = self._api_client.download_file( path, http_options=http_options, + destination=destination, ) - if isinstance(file, types.Video): - file.video_bytes = data - elif isinstance(file, types.GeneratedVideo) and file.video is not None: - file.video.video_bytes = data + if destination is None: + if isinstance(file, types.Video): + file.video_bytes = data + elif isinstance(file, types.GeneratedVideo) and file.video is not None: + file.video.video_bytes = data return data @@ -1277,13 +1310,36 @@ async def upload( kwargs=config_model.model_dump() if config else {}, ) + # from typing import overload + + @overload async def download( self, *, - file: Union[str, types.File], + file: Union[str, types.File, types.Video, types.GeneratedVideo], + destination: None = None, config: Optional[types.DownloadFileConfigOrDict] = None, ) -> bytes: - """Downloads a file's data from the file service. + ... + + @overload + async def download( + self, + *, + file: Union[str, types.File, types.Video, types.GeneratedVideo], + destination: Union[str, os.PathLike[str], io.IOBase], + config: Optional[types.DownloadFileConfigOrDict] = None, + ) -> None: + ... + + async def download( + self, + *, + file: Union[str, types.File, types.Video, types.GeneratedVideo], + destination: Optional[Union[str, os.PathLike[str], io.IOBase]] = None, + config: Optional[types.DownloadFileConfigOrDict] = None, + ) -> Optional[bytes]: + """Downloads a file's data from the file service asynchronously. The Gemini Enterprise Agent Platform implementation of the API does not include the file service. @@ -1291,27 +1347,38 @@ async def download( Files created by `upload` can't be downloaded. You can tell which files are downloadable by checking the `download_uri` property. + Note: When destination is None, this method returns the data as bytes. For + `Video` and `GeneratedVideo` objects, if destination is None, it also sets + the + `video_bytes` property on the `Video` object. When destination is provided, + chunks are streamed directly to disk or the writable stream without loading + the full file into memory, and None is returned. + Args: - File (str): A file name, uri, or file object. Identifying which file to - download. - config (DownloadFileConfigOrDict): Optional, configuration for the get - method. + file: A file name, uri, or file object identifying which file to download. + destination: Optional, a local file path or writable binary stream to + write the file directly to in chunks, avoiding holding the full file in + memory. + config: Optional, configuration for the get method. Returns: - File: The file data as bytes. + bytes: The file data as bytes if destination is None; otherwise None. Usage: .. code-block:: python - for file client.files.list(): + async for file in await client.aio.files.list(): if file.download_uri is not None: break else: raise ValueError('No files found with a `download_uri`.') - data = client.files.download(file=file) - # data = client.files.download(file=file.name) - # data = client.files.download(file=file.uri) + + # Download to memory: + data = await client.aio.files.download(file=file) + + # Stream directly to a file path (avoids memory spike): + await client.aio.files.download(file=file, destination="output.mp4") """ if self._api_client.vertexai: raise ValueError( @@ -1325,6 +1392,12 @@ async def download( else: config_model = config + if isinstance(file, types.File) and file.download_uri is None: + raise ValueError( + "Only generated files can be downloaded, uploaded files can't be " + 'downloaded. You can tell which files are downloadable by checking ' + 'the `source` or `download_uri` property.' + ) name = t.t_file_name(file) path = f'files/{name}:download' @@ -1340,8 +1413,15 @@ async def download( data = await self._api_client.async_download_file( path, http_options=http_options, + destination=destination, ) + if destination is None: + if isinstance(file, types.Video): + file.video_bytes = data + elif isinstance(file, types.GeneratedVideo) and file.video is not None: + file.video.video_bytes = data + return data async def register_files( diff --git a/google/genai/tests/files/test_download.py b/google/genai/tests/files/test_download.py index 1a1627e4f..fa7d030cf 100644 --- a/google/genai/tests/files/test_download.py +++ b/google/genai/tests/files/test_download.py @@ -17,12 +17,28 @@ """Test files upload method.""" +import asyncio +import io import pathlib +import time +from unittest import mock +from typing import Any +import google.auth.transport.requests +import httpx import pytest +import requests from ... import _transformers as t from ... import types +from ..._api_client import BaseApiClient +from ...files import AsyncFiles, Files from .. import pytest_helper +try: + import aiohttp + AIOHTTP_NOT_INSTALLED = False +except ImportError: + AIOHTTP_NOT_INSTALLED = True + test_table: list[pytest_helper.TestTableItem] = [] @@ -36,37 +52,104 @@ pytest_plugins = ('pytest_asyncio',) +def _get_downloadable_file(client: Any) -> types.File | types.Video: + for file in client.files.list(): + if file.download_uri is not None: + return file + # Fallback to generating a minimal video if no downloadable files exist in the project. + operation = client.models.generate_videos( + model='veo-2.0-generate-001', + prompt='A red ball', + config=types.GenerateVideosConfig( + person_generation='dont_allow', + aspect_ratio='16:9', + duration_seconds=5, + ), + ) + while not operation.done: + time.sleep(10) + operation = client.operations.get(operation) + return operation.result.generated_videos[0].video + + +async def _async_get_downloadable_file( + client: Any, +) -> types.File | types.Video: + async for file in await client.aio.files.list(): + if file.download_uri is not None: + return file + # Fallback to generating a minimal video if no downloadable files exist in the project. + operation = await client.aio.models.generate_videos( + model='veo-2.0-generate-001', + prompt='A red ball', + config=types.GenerateVideosConfig( + person_generation='dont_allow', + aspect_ratio='16:9', + duration_seconds=5, + ), + ) + while not operation.done: + await asyncio.sleep(10) + operation = await client.aio.operations.get(operation) + return operation.result.generated_videos[0].video + + +async def _mock_async_download_video(*args: Any, **kwargs: Any) -> bytes | None: + if kwargs.get('destination') is None: + return b'video_data' + return None + + +class _AsyncMockChunkIter: + + async def iter_chunked(self, chunk_size: int | None = None): + yield b'chunk1' + yield b'chunk2' + + +class _AsyncHttpxMockChunkIter: + + async def __call__(self, chunk_size: int | None = None): + yield b'chunk1' + yield b'chunk2' + + +class _AsyncWriter: + + def __init__(self) -> None: + self.written_chunks: list[bytes] = [] + + async def write(self, data: bytes) -> None: + self.written_chunks.append(data) + + def test_name_transform_name(client): with pytest_helper.exception_if_vertex(client, ValueError): - for file in client.files.list(): - if file.download_uri is not None: - break - else: - raise ValueError('No files found with a `download_uri`.') - - file_id = file.name.split('/')[-1] - video = types.Video(uri=file.download_uri) + file = _get_downloadable_file(client) + + file_id = (file.name or file.uri).split('/')[-1].split(':')[0] + download_uri = getattr(file, 'download_uri', None) or getattr( + file, 'uri', None + ) + video = types.Video(uri=download_uri) generated_video = types.GeneratedVideo(video=video) for f in [ file, file_id, - file.name, - file.uri, - file.download_uri, + getattr(file, 'name', file_id), + getattr(file, 'uri', None), + getattr(file, 'download_uri', None), video, generated_video, ]: - name = t.t_file_name(f) - assert name == file_id + if f is not None: + name = t.t_file_name(f) + assert name == file_id def test_basic_download(client): with pytest_helper.exception_if_vertex(client, ValueError): - for file in client.files.list(): - if file.download_uri is not None: - break - else: - raise ValueError('No files found with a `download_uri`.') + file = _get_downloadable_file(client) content = client.files.download(file=file) assert content[4:8] == b'ftyp' @@ -75,11 +158,378 @@ def test_basic_download(client): @pytest.mark.asyncio async def test_basic_download_async(client): with pytest_helper.exception_if_vertex(client, ValueError): - async for file in await client.aio.files.list(): - if file.download_uri is not None: - break - else: - raise ValueError('No files found with a `download_uri`.') + file = await _async_get_downloadable_file(client) content = await client.aio.files.download(file=file) assert content[4:8] == b'ftyp' + + +def test_destination_download(client, tmp_path): + with pytest_helper.exception_if_vertex(client, ValueError): + file = _get_downloadable_file(client) + + out_file = tmp_path / 'downloaded.mp4' + result = client.files.download(file=file, destination=out_file) + assert result is None + assert out_file.exists() + assert out_file.read_bytes()[4:8] == b'ftyp' + + +@pytest.mark.asyncio +async def test_async_destination_download(client, tmp_path): + with pytest_helper.exception_if_vertex(client, ValueError): + file = await _async_get_downloadable_file(client) + + out_file = tmp_path / 'downloaded_async.mp4' + result = await client.aio.files.download(file=file, destination=out_file) + assert result is None + assert out_file.exists() + assert out_file.read_bytes()[4:8] == b'ftyp' + + +def test_destination_filepath(client, tmp_path): + if client._api_client.vertexai: + with pytest.raises( + ValueError, match='only supported in the Gemini Developer client' + ): + client.files.download( + file='files/test_123', destination=str(tmp_path / 'out.mp4') + ) + return + + api_client = mock.create_autospec( + BaseApiClient, instance=True, spec_set=True + ) + api_client.vertexai = False + api_client.download_file.return_value = None + + files_client = Files(api_client) + target_file = str(tmp_path / 'out.mp4') + + result = files_client.download(file='files/test_123', destination=target_file) + assert result is None + api_client.download_file.assert_called_once() + assert api_client.download_file.call_args.kwargs['destination'] == target_file + + +def test_destination_pathlib(client, tmp_path): + if client._api_client.vertexai: + return + + api_client = mock.create_autospec( + BaseApiClient, instance=True, spec_set=True + ) + api_client.vertexai = False + api_client.download_file.return_value = None + + files_client = Files(api_client) + target_file = tmp_path / 'out.mp4' + + result = files_client.download(file='files/test_123', destination=target_file) + assert result is None + api_client.download_file.assert_called_once() + assert api_client.download_file.call_args.kwargs['destination'] == target_file + + +def test_destination_bytesio(client): + if client._api_client.vertexai: + return + + api_client = mock.create_autospec( + BaseApiClient, instance=True, spec_set=True + ) + api_client.vertexai = False + api_client.download_file.return_value = None + + files_client = Files(api_client) + buffer = io.BytesIO() + + result = files_client.download(file='files/test_123', destination=buffer) + assert result is None + api_client.download_file.assert_called_once() + assert api_client.download_file.call_args.kwargs['destination'] == buffer + + +def test_video_destination_behavior(client, tmp_path): + if client._api_client.vertexai: + return + + api_client = mock.create_autospec( + BaseApiClient, instance=True, spec_set=True + ) + api_client.vertexai = False + + # When destination is None, returns bytes and sets video.video_bytes + api_client.download_file.return_value = b'video_data' + files_client = Files(api_client) + video = types.Video( + uri='https://generativelanguage.googleapis.com/v1beta/files/test_video' + ) + data = files_client.download(file=video) + assert data == b'video_data' + assert video.video_bytes == b'video_data' + + # When destination is provided, returns None and does not overwrite video.video_bytes + video2 = types.Video( + uri='https://generativelanguage.googleapis.com/v1beta/files/test_video' + ) + api_client.download_file.return_value = None + buf = io.BytesIO() + data2 = files_client.download(file=video2, destination=buf) + assert data2 is None + assert video2.video_bytes is None + + # GeneratedVideo when destination is None sets video.video_bytes + video3 = types.Video( + uri='https://generativelanguage.googleapis.com/v1beta/files/test_video' + ) + gen_video = types.GeneratedVideo(video=video3) + api_client.download_file.return_value = b'video_data' + data3 = files_client.download(file=gen_video) + assert data3 == b'video_data' + assert gen_video.video.video_bytes == b'video_data' + + +@pytest.mark.asyncio +async def test_async_video_destination_behavior(client): + if client._api_client.vertexai: + return + + api_client = mock.create_autospec( + BaseApiClient, instance=True, spec_set=True + ) + api_client.vertexai = False + api_client.async_download_file = mock.AsyncMock( + side_effect=_mock_async_download_video + ) + files_client = AsyncFiles(api_client) + video = types.Video( + uri='https://generativelanguage.googleapis.com/v1beta/files/test_video' + ) + data = await files_client.download(file=video) + assert data == b'video_data' + assert video.video_bytes == b'video_data' + + video2 = types.Video( + uri='https://generativelanguage.googleapis.com/v1beta/files/test_video' + ) + buf = io.BytesIO() + data2 = await files_client.download(file=video2, destination=buf) + assert data2 is None + assert video2.video_bytes is None + + +@pytest.mark.asyncio +async def test_async_destination(client, tmp_path): + if client._api_client.vertexai: + with pytest.raises( + ValueError, match='only supported in the Gemini Developer client' + ): + await client.aio.files.download( + file='files/test_123', destination=str(tmp_path / 'out.mp4') + ) + return + + api_client = mock.create_autospec( + BaseApiClient, instance=True, spec_set=True + ) + api_client.vertexai = False + api_client.async_download_file = mock.AsyncMock(return_value=None) + files_client = AsyncFiles(api_client) + target_file = tmp_path / 'out.mp4' + + result = await files_client.download( + file='files/test_123', destination=target_file + ) + assert result is None + + +@pytest.mark.asyncio +async def test_async_destination_bytesio(client): + if client._api_client.vertexai: + return + + api_client = mock.create_autospec( + BaseApiClient, instance=True, spec_set=True + ) + api_client.vertexai = False + api_client.async_download_file = mock.AsyncMock(return_value=None) + files_client = AsyncFiles(api_client) + buffer = io.BytesIO() + + result = await files_client.download(file='files/test_123', destination=buffer) + assert result is None + + +def test_destination_invalid_type(client): + if client._api_client.vertexai: + return + + api_client = BaseApiClient(api_key='test_key') + with pytest.raises(ValueError, match='Unsupported destination type'): + api_client.download_file('files/test_123:download', destination=12345) + + +@pytest.mark.asyncio +async def test_async_destination_invalid_type(client): + if client._api_client.vertexai: + return + + api_client = BaseApiClient(api_key='test_key') + with pytest.raises(ValueError, match='Unsupported destination type'): + await api_client.async_download_file( + 'files/test_123:download', destination=12345 + ) + + +@pytest.mark.asyncio +@pytest.mark.skipif( + AIOHTTP_NOT_INSTALLED, reason='aiohttp is not installed, skipping test.' +) +async def test_async_destination_bytesio_writes_chunks(client): + if client._api_client.vertexai: + return + + api_client = BaseApiClient(api_key='test_key') + + mock_response = mock.create_autospec(aiohttp.ClientResponse, instance=True) + mock_response.status = 200 + mock_response.content = _AsyncMockChunkIter() + + mock_session = mock.create_autospec(aiohttp.ClientSession, instance=True) + mock_session.request = mock.AsyncMock(return_value=mock_response) + mock_session.configure_mtls_channel = mock.AsyncMock() + mock_session._is_mtls = False + + buffer = io.BytesIO() + with mock.patch.object( + api_client, '_use_aiohttp', return_value=True + ), mock.patch.object( + api_client, '_get_aiohttp_session', return_value=mock_session + ): + result = await api_client.async_download_file( + 'files/test_123:download', destination=buffer + ) + + assert result is None + assert buffer.getvalue() == b'chunk1chunk2' + mock_response.close.assert_called_once() + + +def test_authorized_session_destination_closes_response(client): + if client._api_client.vertexai: + return + + api_client = BaseApiClient(api_key='test_key') + mock_response = mock.create_autospec(requests.Response, instance=True) + mock_response.status_code = 200 + mock_response.iter_content.return_value = [b'chunk1', b'chunk2'] + + mock_auth_session = mock.create_autospec( + google.auth.transport.requests.AuthorizedSession, instance=True + ) + mock_auth_session.request.return_value = mock_response + mock_auth_session._is_mtls = False + + api_client._authorized_session = mock_auth_session + buffer = io.BytesIO() + + with mock.patch.object(api_client, '_use_google_auth_sync', return_value=True): + result = api_client.download_file( + 'files/test_123:download', destination=buffer + ) + + assert result is None + assert buffer.getvalue() == b'chunk1chunk2' + mock_response.close.assert_called_once() + + +@pytest.mark.asyncio +@pytest.mark.skipif( + AIOHTTP_NOT_INSTALLED, reason='aiohttp is not installed, skipping test.' +) +async def test_async_destination_awaitable_writer(client): + if client._api_client.vertexai: + return + + api_client = BaseApiClient(api_key='test_key') + + mock_response = mock.create_autospec(aiohttp.ClientResponse, instance=True) + mock_response.status = 200 + mock_response.content = _AsyncMockChunkIter() + + mock_session = mock.create_autospec(aiohttp.ClientSession, instance=True) + mock_session.request = mock.AsyncMock(return_value=mock_response) + mock_session.configure_mtls_channel = mock.AsyncMock() + mock_session._is_mtls = False + + async_writer = _AsyncWriter() + with mock.patch.object( + api_client, '_use_aiohttp', return_value=True + ), mock.patch.object( + api_client, '_get_aiohttp_session', return_value=mock_session + ): + result = await api_client.async_download_file( + 'files/test_123:download', destination=async_writer + ) + + assert result is None + assert async_writer.written_chunks == [b'chunk1', b'chunk2'] + mock_response.close.assert_called_once() + + +@pytest.mark.asyncio +async def test_async_httpx_destination_bytesio_writes_chunks(client): + if client._api_client.vertexai: + return + + api_client = BaseApiClient(api_key='test_key') + + mock_response = mock.create_autospec(httpx.Response, instance=True) + mock_response.status_code = 200 + mock_response.aiter_bytes = _AsyncHttpxMockChunkIter() + mock_response.aclose = mock.AsyncMock() + + buffer = io.BytesIO() + with mock.patch.object( + api_client, '_use_aiohttp', return_value=False + ), mock.patch.object( + api_client._async_httpx_client, 'send', mock.AsyncMock(return_value=mock_response) + ): + result = await api_client.async_download_file( + 'files/test_123:download', destination=buffer + ) + + assert result is None + assert buffer.getvalue() == b'chunk1chunk2' + mock_response.aclose.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_async_httpx_destination_awaitable_writer(client): + if client._api_client.vertexai: + return + + api_client = BaseApiClient(api_key='test_key') + + mock_response = mock.create_autospec(httpx.Response, instance=True) + mock_response.status_code = 200 + mock_response.aiter_bytes = _AsyncHttpxMockChunkIter() + mock_response.aclose = mock.AsyncMock() + + async_writer = _AsyncWriter() + with mock.patch.object( + api_client, '_use_aiohttp', return_value=False + ), mock.patch.object( + api_client._async_httpx_client, 'send', mock.AsyncMock(return_value=mock_response) + ): + result = await api_client.async_download_file( + 'files/test_123:download', destination=async_writer + ) + + assert result is None + assert async_writer.written_chunks == [b'chunk1', b'chunk2'] + mock_response.aclose.assert_awaited_once() + + +