diff --git a/google/genai/_gaos/google_genai.py b/google/genai/_gaos/google_genai.py index dc9b3e44c..f9bc83e7b 100644 --- a/google/genai/_gaos/google_genai.py +++ b/google/genai/_gaos/google_genai.py @@ -45,6 +45,7 @@ wrap_stream_errors, ) from .sdk import AsyncGenAI, GenAI +from .types import environments from .types import interactions from .types.security import Security from .utils import BackoffStrategy, RetryConfig, eventstreaming @@ -761,10 +762,69 @@ async def list_executions(self, *args: Any, **kwargs: Any) -> Any: return await async_wrap_sdk_call(super().list_executions, *args, **kwargs) +class GeminiNextGenEnvironmentFiles: + """Environment files resource backed by the NextGen client.""" + + def __init__(self, parent: GeminiNextGenEnvironments): + self._parent = parent + + def get(self, *args: Any, **kwargs: Any) -> Any: + files_sdk = getattr(super(GeminiNextGenEnvironments, self._parent), 'files', None) + if files_sdk is not None and hasattr(files_sdk, 'get'): + return wrap_sdk_call(files_sdk.get, *args, **kwargs) + raise AttributeError( + 'environments.files.get is only available in the private SDK.' + ) + + def download( + self, + *, + environment: str, + path: str, + http_options: Optional[Any] = None, + ) -> bytes: + """Downloads binary file content from an environment workspace.""" + return self._parent.download_file( + environment=environment, + path=path, + http_options=http_options, + ) + + +class AsyncGeminiNextGenEnvironmentFiles: + """Async environment files resource backed by the NextGen client.""" + + def __init__(self, parent: AsyncGeminiNextGenEnvironments): + self._parent = parent + + async def get(self, *args: Any, **kwargs: Any) -> Any: + files_sdk = getattr(super(AsyncGeminiNextGenEnvironments, self._parent), 'files', None) + if files_sdk is not None and hasattr(files_sdk, 'get'): + return await async_wrap_sdk_call(files_sdk.get, *args, **kwargs) + raise AttributeError( + 'environments.files.get is only available in the private SDK.' + ) + + async def download( + self, + *, + environment: str, + path: str, + http_options: Optional[Any] = None, + ) -> bytes: + """Downloads binary file content from an environment workspace.""" + return await self._parent.download_file( + environment=environment, + path=path, + http_options=http_options, + ) + + class GeminiNextGenEnvironments(GeneratedEnvironments): """Public environments resource backed by the NextGen client.""" def __init__(self, api_client: Any): + self._api_client = api_client sdk = build_google_genai_client(api_client) super().__init__(sdk.sdk_configuration, parent_ref=sdk) @@ -777,6 +837,10 @@ def with_raw_response(self): def with_streaming_response(self): return _RawResponseAccessorProxy(super().with_streaming_response) + @property + def files(self) -> GeminiNextGenEnvironmentFiles: + return GeminiNextGenEnvironmentFiles(self) + def create_environment(self, *args: Any, **kwargs: Any) -> Any: return wrap_sdk_call(super().create_environment, *args, **kwargs) @@ -802,15 +866,37 @@ def delete(self, *args: Any, **kwargs: Any) -> Any: return self.delete_environment(*args, **kwargs) def get_environment_files(self, *args: Any, **kwargs: Any) -> Any: - return wrap_sdk_call(super().get_environment_files, *args, **kwargs) + return self.files.get(*args, **kwargs) + + def get_files(self, *args: Any, **kwargs: Any) -> Any: + return self.files.get(*args, **kwargs) - # NOTE: update_environment, patch_environment are handled by fallback if they exist, but we assume they aren't generated based on our openapi.json. + def download_file( + self, + *, + environment: str, + path: str, + http_options: Optional[Any] = None, + ) -> bytes: + """Downloads binary file content from an environment workspace.""" + env_name = ( + environment + if environment.startswith('environments/') + else f'environments/{environment}' + ) + clean_path = path.lstrip('/') + download_path = f'{env_name}/files/{clean_path}?alt=media' + return self._api_client.download_file( + download_path, + http_options=http_options, + ) class AsyncGeminiNextGenEnvironments(GeneratedAsyncEnvironments): """Async public environments resource backed by the NextGen client.""" def __init__(self, api_client: Any): + self._api_client = api_client sdk = build_google_genai_async_client(api_client) super().__init__(sdk.sdk_configuration, parent_ref=sdk) @@ -823,6 +909,10 @@ def with_raw_response(self): def with_streaming_response(self): return _AsyncRawResponseAccessorProxy(super().with_streaming_response) + @property + def files(self) -> AsyncGeminiNextGenEnvironmentFiles: + return AsyncGeminiNextGenEnvironmentFiles(self) + async def create_environment(self, *args: Any, **kwargs: Any) -> Any: return await async_wrap_sdk_call(super().create_environment, *args, **kwargs) @@ -848,7 +938,30 @@ async def delete(self, *args: Any, **kwargs: Any) -> Any: return await self.delete_environment(*args, **kwargs) async def get_environment_files(self, *args: Any, **kwargs: Any) -> Any: - return await async_wrap_sdk_call(super().get_environment_files, *args, **kwargs) + return await self.files.get(*args, **kwargs) + + async def get_files(self, *args: Any, **kwargs: Any) -> Any: + return await self.files.get(*args, **kwargs) + + async def download_file( + self, + *, + environment: str, + path: str, + http_options: Optional[Any] = None, + ) -> bytes: + """Downloads binary file content from an environment workspace.""" + env_name = ( + environment + if environment.startswith('environments/') + else f'environments/{environment}' + ) + clean_path = path.lstrip('/') + download_path = f'{env_name}/files/{clean_path}?alt=media' + return await self._api_client.async_download_file( + download_path, + http_options=http_options, + ) def _add_output_properties_if_interaction(value: Any) -> Any: diff --git a/google/genai/tests/gaos/test_environments_lifecycle.py b/google/genai/tests/gaos/test_environments_lifecycle.py index bc1fdb6b6..894ca4cb4 100644 --- a/google/genai/tests/gaos/test_environments_lifecycle.py +++ b/google/genai/tests/gaos/test_environments_lifecycle.py @@ -116,3 +116,145 @@ def test_python_environments_lifecycle_routes_through_google_genai_client( server.shutdown() thread.join() server.server_close() + + +class _ScottyDownloadHandler(BaseHTTPRequestHandler): + captured: list[str] = [] + + def do_GET(self) -> None: + self.captured.append(f"GET {self.path}") + if "?alt=media" in self.path: + payload = b"print('downloaded content')\n" + self.send_response(200) + self.send_header("content-type", "application/octet-stream") + self.send_header("content-length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + return + elif "/files" in self.path: + payload = json.dumps({"files": [{"name": "main.py"}]}).encode() + self.send_response(200) + self.send_header("content-type", "application/json") + self.send_header("content-length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + return + + self.send_response(404) + self.end_headers() + + def log_message(self, *args) -> None: + pass + + +def test_python_environments_files_get_and_download(monkeypatch): + monkeypatch.delenv("GOOGLE_GENAI_USE_VERTEXAI", raising=False) + captured: list[str] = [] + handler = type("Handler", (_ScottyDownloadHandler,), { + "captured": captured, + }) + server = ThreadingHTTPServer(("127.0.0.1", 0), handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + client = Client( + api_key="test-api-key", + http_options={ + "api_version": "v1beta", + "base_url": f"http://127.0.0.1:{server.server_port}", + "headers": {"X-Goog-Api-Client": "test"}, + }, + ) + + # Test sync files.get + if getattr(client.environments, "_sdk_files", None) is not None: + files_res = client.environments.files.get( + environment="env_123", + path="src/main.py", + ) + files_list = files_res.files if hasattr(files_res, "files") else files_res.get("files", []) + assert len(files_list) == 1 + assert (files_list[0].name if hasattr(files_list[0], "name") else files_list[0]["name"]) == "main.py" + else: + with pytest.raises(AttributeError, match="only available in the private SDK"): + client.environments.files.get( + environment="env_123", + path="src/main.py", + ) + + # Test sync files.download + downloaded = client.environments.files.download( + environment="env_123", + path="src/main.py", + ) + assert downloaded == b"print('downloaded content')\n" + + # Test download_file convenience alias + downloaded_alias = client.environments.download_file( + environment="env_123", + path="src/main.py", + ) + assert downloaded_alias == b"print('downloaded content')\n" + + finally: + server.shutdown() + thread.join() + server.server_close() + + +@pytest.mark.asyncio +async def test_python_environments_async_files_get_and_download(monkeypatch): + monkeypatch.delenv("GOOGLE_GENAI_USE_VERTEXAI", raising=False) + captured: list[str] = [] + handler = type("Handler", (_ScottyDownloadHandler,), { + "captured": captured, + }) + server = ThreadingHTTPServer(("127.0.0.1", 0), handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + client = Client( + api_key="test-api-key", + http_options={ + "api_version": "v1beta", + "base_url": f"http://127.0.0.1:{server.server_port}", + }, + ) + + # Test async files.get + if getattr(client.aio.environments, "_sdk_files", None) is not None: + files_res = await client.aio.environments.files.get( + environment="env_123", + path="src/main.py", + ) + files_list = files_res.files if hasattr(files_res, "files") else files_res.get("files", []) + assert len(files_list) == 1 + assert (files_list[0].name if hasattr(files_list[0], "name") else files_list[0]["name"]) == "main.py" + else: + with pytest.raises(AttributeError, match="only available in the private SDK"): + await client.aio.environments.files.get( + environment="env_123", + path="src/main.py", + ) + + # Test async files.download + downloaded = await client.aio.environments.files.download( + environment="env_123", + path="src/main.py", + ) + assert downloaded == b"print('downloaded content')\n" + + # Test async download_file convenience alias + downloaded_alias = await client.aio.environments.download_file( + environment="env_123", + path="src/main.py", + ) + assert downloaded_alias == b"print('downloaded content')\n" + + finally: + server.shutdown() + thread.join() + server.server_close() + + +