From b56217454106af63ede07dd9f16c709ba6bf0dba Mon Sep 17 00:00:00 2001 From: Mark Daoust Date: Wed, 12 Aug 2026 16:03:10 -0700 Subject: [PATCH] feat: implement environments files upload and download across Go, Java, .NET, Python, and TypeScript SDKs PiperOrigin-RevId: 963709775 --- google/genai/_gaos/google_genai.py | 318 +++++++++++++++++- .../tests/gaos/test_environments_lifecycle.py | 85 ++++- 2 files changed, 390 insertions(+), 13 deletions(-) diff --git a/google/genai/_gaos/google_genai.py b/google/genai/_gaos/google_genai.py index d1d83e84a..c83548272 100644 --- a/google/genai/_gaos/google_genai.py +++ b/google/genai/_gaos/google_genai.py @@ -26,6 +26,11 @@ from typing import TYPE_CHECKING, Any, Mapping, Optional, TypeVar, Union, cast +import io +import json +import mimetypes +import os + import httpx from ._hooks.google_genai_auth import ( @@ -765,6 +770,7 @@ async def list_executions(self, *args: Any, **kwargs: Any) -> Any: return await async_wrap_sdk_call(super().list_executions, *args, **kwargs) + class GeminiNextGenEnvironmentFiles(GeneratedFiles): """Environment files resource backed by the NextGen client.""" @@ -793,17 +799,21 @@ def list(self, *args: Any, **kwargs: Any) -> Any: def download( self, *, - environment: str, path: str, + environment: Optional[str] = None, + environment_id: Optional[str] = None, http_options: Optional[Any] = None, ) -> bytes: """Downloads binary file content from an environment workspace.""" if not self._api_client: raise AttributeError('api_client is required to download files.') + target_env = environment or environment_id + if not target_env: + raise ValueError('environment or environment_id is required.') env_name = ( - environment - if environment.startswith('environments/') - else f'environments/{environment}' + target_env + if target_env.startswith('environments/') + else f'environments/{target_env}' ) clean_path = path.lstrip('/') download_path = f'{env_name}/files/{clean_path}?alt=media' @@ -812,6 +822,150 @@ def download( http_options=http_options, ) + def upload( + self, + *, + path: str, + file: Union[str, os.PathLike[str], io.IOBase, bytes], + environment: Optional[str] = None, + environment_id: Optional[str] = None, + mime_type: Optional[str] = None, + overwrite: Optional[bool] = None, + extract: Optional[bool] = None, + http_options: Optional[Any] = None, + ) -> Union[environments.GetEnvironmentFilesResponse, Any]: + """Uploads a file or extracts an archive inside an environment workspace.""" + if not self._api_client: + raise AttributeError('api_client is required to upload files.') + target_env = environment or environment_id + if not target_env: + raise ValueError('environment or environment_id is required.') + env_name = ( + target_env + if target_env.startswith('environments/') + else f'environments/{target_env}' + ) + clean_path = path.lstrip('/') + + file_obj: Union[str, io.IOBase] + if isinstance(file, (bytes, bytearray)): + file_obj = io.BytesIO(file) + size_bytes = len(file) + elif isinstance(file, io.IOBase): + file_obj = file + offset = file_obj.tell() + file_obj.seek(0, os.SEEK_END) + size_bytes = file_obj.tell() - offset + file_obj.seek(offset, os.SEEK_SET) + else: + fs_path = os.fspath(file) + if not fs_path or not os.path.isfile(fs_path): + raise FileNotFoundError(f'{file} is not a valid file path.') + size_bytes = os.path.getsize(fs_path) + file_obj = fs_path + if mime_type is None: + mime_type, _ = mimetypes.guess_type(fs_path) + + if mime_type is None: + mime_type = 'application/octet-stream' + + query_params = [] + if overwrite is not None: + query_params.append(f'overwrite={"true" if overwrite else "false"}') + if extract is not None: + query_params.append(f'extract={"true" if extract else "false"}') + query_str = '&'.join(query_params) + handshake_path = f'{env_name}/files/{clean_path}' + if query_str: + handshake_path = f'{handshake_path}?{query_str}' + + user_headers = {} + if http_options: + if isinstance(http_options, dict): + user_headers = http_options.get('headers', {}) or {} + elif hasattr(http_options, 'headers') and http_options.headers: + user_headers = dict(http_options.headers) + + upload_headers = { + **user_headers, + 'X-Goog-Upload-Protocol': 'resumable', + 'X-Goog-Upload-Command': 'start', + 'X-Goog-Upload-Header-Content-Length': str(size_bytes), + 'X-Goog-Upload-Header-Content-Type': mime_type, + } + + req_api_version = ( + ( + http_options.get('api_version') + if isinstance(http_options, dict) + else getattr(http_options, 'api_version', None) + ) + or self._api_client._http_options.api_version + or 'v1alpha' + ).lstrip('/') + + if http_options: + if isinstance(http_options, dict): + merged_options = { + **http_options, + 'headers': upload_headers, + 'api_version': f'upload/{req_api_version}', + } + else: + merged_options = http_options.model_copy() + merged_options.headers = upload_headers + merged_options.api_version = f'upload/{req_api_version}' + else: + merged_options = { + 'headers': upload_headers, + 'api_version': f'upload/{req_api_version}', + } + + response = self._api_client.request( + 'put', + handshake_path, + request_dict={}, + http_options=merged_options, + ) + + if ( + response is None + or response.headers is None + or 'x-goog-upload-url' not in response.headers + ): + raise KeyError( + 'Failed to upload file: Upload URL was not returned from the upload request.' + ) + upload_url = response.headers['x-goog-upload-url'] + + upload_response = self._api_client.upload_file( + file_obj, + upload_url, + size_bytes, + http_options=http_options, + ) + + body_text = ( + upload_response.response_stream[0] + if upload_response and upload_response.response_stream + else '{}' + ) + try: + res_json = json.loads(body_text) if body_text else {} + except Exception: + return body_text + + if isinstance(res_json, dict): + if 'files' in res_json and isinstance(res_json['files'], list): + return environments.GetEnvironmentFilesResponse.model_validate(res_json) + elif 'name' in res_json or 'path' in res_json: + file_obj = environments.EnvironmentFile.model_validate(res_json) + return environments.GetEnvironmentFilesResponse(files=[file_obj]) + elif 'file' in res_json and isinstance(res_json['file'], dict): + file_obj = environments.EnvironmentFile.model_validate(res_json['file']) + return environments.GetEnvironmentFilesResponse(files=[file_obj]) + return res_json + class AsyncGeminiNextGenEnvironmentFiles(GeneratedAsyncFiles): """Async environment files resource backed by the NextGen client.""" @@ -841,17 +995,21 @@ async def list(self, *args: Any, **kwargs: Any) -> Any: async def download( self, *, - environment: str, path: str, + environment: Optional[str] = None, + environment_id: Optional[str] = None, http_options: Optional[Any] = None, ) -> bytes: """Downloads binary file content from an environment workspace.""" if not self._api_client: raise AttributeError('api_client is required to download files.') + target_env = environment or environment_id + if not target_env: + raise ValueError('environment or environment_id is required.') env_name = ( - environment - if environment.startswith('environments/') - else f'environments/{environment}' + target_env + if target_env.startswith('environments/') + else f'environments/{target_env}' ) clean_path = path.lstrip('/') download_path = f'{env_name}/files/{clean_path}?alt=media' @@ -860,6 +1018,150 @@ async def download( http_options=http_options, ) + async def upload( + self, + *, + path: str, + file: Union[str, os.PathLike[str], io.IOBase, bytes], + environment: Optional[str] = None, + environment_id: Optional[str] = None, + mime_type: Optional[str] = None, + overwrite: Optional[bool] = None, + extract: Optional[bool] = None, + http_options: Optional[Any] = None, + ) -> Union[environments.GetEnvironmentFilesResponse, Any]: + """Uploads a file or extracts an archive inside an environment workspace.""" + if not self._api_client: + raise AttributeError('api_client is required to upload files.') + target_env = environment or environment_id + if not target_env: + raise ValueError('environment or environment_id is required.') + env_name = ( + target_env + if target_env.startswith('environments/') + else f'environments/{target_env}' + ) + clean_path = path.lstrip('/') + + file_obj: Union[str, io.IOBase] + if isinstance(file, (bytes, bytearray)): + file_obj = io.BytesIO(file) + size_bytes = len(file) + elif isinstance(file, io.IOBase): + file_obj = file + offset = file_obj.tell() + file_obj.seek(0, os.SEEK_END) + size_bytes = file_obj.tell() - offset + file_obj.seek(offset, os.SEEK_SET) + else: + fs_path = os.fspath(file) + if not fs_path or not os.path.isfile(fs_path): + raise FileNotFoundError(f'{file} is not a valid file path.') + size_bytes = os.path.getsize(fs_path) + file_obj = fs_path + if mime_type is None: + mime_type, _ = mimetypes.guess_type(fs_path) + + if mime_type is None: + mime_type = 'application/octet-stream' + + query_params = [] + if overwrite is not None: + query_params.append(f'overwrite={"true" if overwrite else "false"}') + if extract is not None: + query_params.append(f'extract={"true" if extract else "false"}') + query_str = '&'.join(query_params) + handshake_path = f'{env_name}/files/{clean_path}' + if query_str: + handshake_path = f'{handshake_path}?{query_str}' + + user_headers = {} + if http_options: + if isinstance(http_options, dict): + user_headers = http_options.get('headers', {}) or {} + elif hasattr(http_options, 'headers') and http_options.headers: + user_headers = dict(http_options.headers) + + upload_headers = { + **user_headers, + 'X-Goog-Upload-Protocol': 'resumable', + 'X-Goog-Upload-Command': 'start', + 'X-Goog-Upload-Header-Content-Length': str(size_bytes), + 'X-Goog-Upload-Header-Content-Type': mime_type, + } + + req_api_version = ( + ( + http_options.get('api_version') + if isinstance(http_options, dict) + else getattr(http_options, 'api_version', None) + ) + or self._api_client._http_options.api_version + or 'v1alpha' + ).lstrip('/') + + if http_options: + if isinstance(http_options, dict): + merged_options = { + **http_options, + 'headers': upload_headers, + 'api_version': f'upload/{req_api_version}', + } + else: + merged_options = http_options.model_copy() + merged_options.headers = upload_headers + merged_options.api_version = f'upload/{req_api_version}' + else: + merged_options = { + 'headers': upload_headers, + 'api_version': f'upload/{req_api_version}', + } + + response = await self._api_client.async_request( + 'put', + handshake_path, + request_dict={}, + http_options=merged_options, + ) + + if ( + response is None + or response.headers is None + or 'x-goog-upload-url' not in response.headers + ): + raise KeyError( + 'Failed to upload file: Upload URL was not returned from the upload request.' + ) + upload_url = response.headers['x-goog-upload-url'] + + upload_response = await self._api_client.async_upload_file( + file_obj, + upload_url, + size_bytes, + http_options=http_options, + ) + + body_text = ( + upload_response.response_stream[0] + if upload_response and upload_response.response_stream + else '{}' + ) + try: + res_json = json.loads(body_text) if body_text else {} + except Exception: + return body_text + + if isinstance(res_json, dict): + if 'files' in res_json and isinstance(res_json['files'], list): + return environments.GetEnvironmentFilesResponse.model_validate(res_json) + elif 'name' in res_json or 'path' in res_json: + file_obj = environments.EnvironmentFile.model_validate(res_json) + return environments.GetEnvironmentFilesResponse(files=[file_obj]) + elif 'file' in res_json and isinstance(res_json['file'], dict): + file_obj = environments.EnvironmentFile.model_validate(res_json['file']) + return environments.GetEnvironmentFilesResponse(files=[file_obj]) + return res_json + class GeminiNextGenEnvironments(GeneratedEnvironments): """Public environments resource backed by the NextGen client.""" diff --git a/google/genai/tests/gaos/test_environments_lifecycle.py b/google/genai/tests/gaos/test_environments_lifecycle.py index 39cb3c90c..c599f965b 100644 --- a/google/genai/tests/gaos/test_environments_lifecycle.py +++ b/google/genai/tests/gaos/test_environments_lifecycle.py @@ -98,6 +98,8 @@ def test_python_environments_lifecycle_routes_through_google_genai_client( monkeypatch, ): monkeypatch.delenv("GOOGLE_GENAI_USE_VERTEXAI", raising=False) + for var in ("http_proxy", "https_proxy", "all_proxy", "HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY"): + monkeypatch.delenv(var, raising=False) captured: list[str] = [] captured_bodies: list[dict] = [] handler = type("Handler", (_RecordingHandler,), { @@ -195,8 +197,49 @@ async def test_python_environments_async_create_with_from_environment( server.server_close() -class _ScottyDownloadHandler(BaseHTTPRequestHandler): +class _ScottyFileHandler(BaseHTTPRequestHandler): captured: list[str] = [] + uploaded_bytes: list[bytes] = [] + + def do_PUT(self) -> None: + self.captured.append(f"PUT {self.path}") + if self.path.startswith("/upload/") and ("/environments/" in self.path) and ("/files/" in self.path): + # Initial Scotty upload handshake + upload_url = f"http://127.0.0.1:{self.server.server_port}/scotty/upload/resumable_123" + self.send_response(200) + self.send_header("x-goog-upload-url", upload_url) + self.send_header("x-goog-upload-status", "active") + self.send_header("content-length", "0") + self.end_headers() + return + + self.send_response(404) + self.end_headers() + + def do_POST(self) -> None: + self.captured.append(f"POST {self.path}") + if self.path == "/scotty/upload/resumable_123": + content_length = int(self.headers.get("Content-Length", 0)) + data = self.rfile.read(content_length) + self.uploaded_bytes.append(data) + file_response = { + "file": { + "name": "main.py", + "sizeBytes": str(len(data)), + "mimeType": "text/x-python", + } + } + payload = json.dumps(file_response).encode() + self.send_response(200) + self.send_header("content-type", "application/json") + self.send_header("x-goog-upload-status", "final") + self.send_header("content-length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + return + + self.send_response(404) + self.end_headers() def do_GET(self) -> None: self.captured.append(f"GET {self.path}") @@ -224,11 +267,15 @@ def log_message(self, *args) -> None: pass -def test_python_environments_files_list_and_download(monkeypatch): +def test_python_environments_file_upload_download(monkeypatch): monkeypatch.delenv("GOOGLE_GENAI_USE_VERTEXAI", raising=False) + for var in ("http_proxy", "https_proxy", "all_proxy", "HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY"): + monkeypatch.delenv(var, raising=False) captured: list[str] = [] - handler = type("Handler", (_ScottyDownloadHandler,), { + uploaded_bytes: list[bytes] = [] + handler = type("Handler", (_ScottyFileHandler,), { "captured": captured, + "uploaded_bytes": uploaded_bytes, }) server = ThreadingHTTPServer(("127.0.0.1", 0), handler) thread = threading.Thread(target=server.serve_forever, daemon=True) @@ -255,6 +302,7 @@ def test_python_environments_files_list_and_download(monkeypatch): assert files_res.files[0].size_bytes == 128 assert files_res.next_page_token == "token_next_123" + # Test sync files.list with pagination and recursive options files_res_paginated = client.environments.files.list( environment="env_123", @@ -265,6 +313,17 @@ def test_python_environments_files_list_and_download(monkeypatch): ) assert len(files_res_paginated.files) == 1 + # Test sync upload + upload_res = client.environments.files.upload( + environment="env_123", + path="src/main.py", + file=b"print('hello world')", + mime_type="text/x-python", + ) + assert upload_res.files and len(upload_res.files) == 1 + assert upload_res.files[0].name == "main.py" + assert uploaded_bytes[0] == b"print('hello world')" + # Test sync files.download downloaded = client.environments.files.download( environment="env_123", @@ -299,11 +358,15 @@ def test_python_environments_files_list_and_download(monkeypatch): @pytest.mark.asyncio -async def test_python_environments_async_files_list_and_download(monkeypatch): +async def test_python_environments_async_file_upload_download(monkeypatch): monkeypatch.delenv("GOOGLE_GENAI_USE_VERTEXAI", raising=False) + for var in ("http_proxy", "https_proxy", "all_proxy", "HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY"): + monkeypatch.delenv(var, raising=False) captured: list[str] = [] - handler = type("Handler", (_ScottyDownloadHandler,), { + uploaded_bytes: list[bytes] = [] + handler = type("Handler", (_ScottyFileHandler,), { "captured": captured, + "uploaded_bytes": uploaded_bytes, }) server = ThreadingHTTPServer(("127.0.0.1", 0), handler) thread = threading.Thread(target=server.serve_forever, daemon=True) @@ -314,6 +377,7 @@ async def test_python_environments_async_files_list_and_download(monkeypatch): http_options={ "api_version": "v1beta", "base_url": f"http://127.0.0.1:{server.server_port}", + "headers": {"X-Goog-Api-Client": "test"}, }, ) @@ -338,6 +402,17 @@ async def test_python_environments_async_files_list_and_download(monkeypatch): ) assert len(files_res_paginated.files) == 1 + # Test async upload + upload_res = await client.aio.environments.files.upload( + environment="env_123", + path="src/main.py", + file=b"print('async hello world')", + mime_type="text/x-python", + ) + assert upload_res.files and len(upload_res.files) == 1 + assert upload_res.files[0].name == "main.py" + assert uploaded_bytes[0] == b"print('async hello world')" + # Test async files.download downloaded = await client.aio.environments.files.download( environment="env_123",