diff --git a/google/genai/_gaos/google_genai.py b/google/genai/_gaos/google_genai.py index d1d83e84a..670d70069 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 ( @@ -793,17 +798,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 +821,130 @@ 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.EnvironmentFile, 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, + } + + if http_options: + if isinstance(http_options, dict): + merged_options = {**http_options, 'headers': upload_headers} + else: + merged_options = http_options.model_copy() + merged_options.headers = upload_headers + else: + merged_options = {'headers': upload_headers} + + 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 'file' in res_json and isinstance(res_json['file'], dict): + return environments.EnvironmentFile.model_validate(res_json['file']) + elif '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: + return environments.EnvironmentFile.model_validate(res_json) + return res_json + class AsyncGeminiNextGenEnvironmentFiles(GeneratedAsyncFiles): """Async environment files resource backed by the NextGen client.""" @@ -841,17 +974,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 +997,130 @@ 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.EnvironmentFile, 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, + } + + if http_options: + if isinstance(http_options, dict): + merged_options = {**http_options, 'headers': upload_headers} + else: + merged_options = http_options.model_copy() + merged_options.headers = upload_headers + else: + merged_options = {'headers': upload_headers} + + 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 'file' in res_json and isinstance(res_json['file'], dict): + return environments.EnvironmentFile.model_validate(res_json['file']) + elif '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: + return environments.EnvironmentFile.model_validate(res_json) + 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 224602593..951aa90c6 100644 --- a/google/genai/tests/gaos/test_environments_lifecycle.py +++ b/google/genai/tests/gaos/test_environments_lifecycle.py @@ -143,8 +143,49 @@ def test_python_environments_lifecycle_routes_through_google_genai_client( 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 ("/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}") @@ -172,11 +213,13 @@ 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) 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) @@ -213,6 +256,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", + ) + name = upload_res.name if hasattr(upload_res, "name") else upload_res.get("name") + assert name == "main.py" + assert uploaded_bytes[0] == b"print('hello world')" + # Test sync files.download downloaded = client.environments.files.download( environment="env_123", @@ -247,11 +301,13 @@ 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) 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) @@ -262,6 +318,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"}, }, ) @@ -286,6 +343,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", + ) + name = upload_res.name if hasattr(upload_res, "name") else upload_res.get("name") + assert 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", @@ -372,6 +440,3 @@ def test_python_environments_types_and_models(): assert req.recursive is True assert req.api_version == "v1beta" - - -