diff --git a/google/genai/_gaos/environments.py b/google/genai/_gaos/environments.py index 48de2a95a..db2e4a41c 100644 --- a/google/genai/_gaos/environments.py +++ b/google/genai/_gaos/environments.py @@ -234,6 +234,7 @@ def create_environment( self, *, api_version: Optional[str] = None, + from_environment: Optional[str] = None, network: Optional[ Union[ environments_createenvironmentrequest.CreateEnvironmentRequestNetworkUnion, @@ -254,6 +255,9 @@ def create_environment( r"""Creates an environment. :param api_version: Which version of the API to use. + :param from_environment: Optional. The source environment to copy/fork from. + Format: `environments/{environment_id}` or `{environment_id}`. + When specified, `sources` and `env` must be empty. :param network: Network configuration for the environment. :param sources: Sources to be mounted into the environment. :param extra_headers: Additional headers to set or replace on requests. @@ -278,6 +282,7 @@ def create_environment( request = models.CreateEnvironmentRequest( api_version=api_version, body=environments.CreateEnvironmentRequest( + from_environment=from_environment, network=utils.get_pydantic_model( network, Optional[environments.CreateEnvironmentRequestNetworkUnion] ), @@ -1025,6 +1030,7 @@ async def create_environment( self, *, api_version: Optional[str] = None, + from_environment: Optional[str] = None, network: Optional[ Union[ environments_createenvironmentrequest.CreateEnvironmentRequestNetworkUnion, @@ -1045,6 +1051,9 @@ async def create_environment( r"""Creates an environment. :param api_version: Which version of the API to use. + :param from_environment: Optional. The source environment to copy/fork from. + Format: `environments/{environment_id}` or `{environment_id}`. + When specified, `sources` and `env` must be empty. :param network: Network configuration for the environment. :param sources: Sources to be mounted into the environment. :param extra_headers: Additional headers to set or replace on requests. @@ -1069,6 +1078,7 @@ async def create_environment( request = models.CreateEnvironmentRequest( api_version=api_version, body=environments.CreateEnvironmentRequest( + from_environment=from_environment, network=utils.get_pydantic_model( network, Optional[environments.CreateEnvironmentRequestNetworkUnion] ), diff --git a/google/genai/_gaos/types/environments/createenvironmentrequest.py b/google/genai/_gaos/types/environments/createenvironmentrequest.py index 59a53b4ba..7b9143204 100644 --- a/google/genai/_gaos/types/environments/createenvironmentrequest.py +++ b/google/genai/_gaos/types/environments/createenvironmentrequest.py @@ -54,6 +54,11 @@ class CreateEnvironmentRequestParam(TypedDict): r"""Request for `CreateEnvironment`.""" + from_environment: NotRequired[str] + r"""Optional. The source environment to copy/fork from. + Format: `environments/{environment_id}` or `{environment_id}`. + When specified, `sources` and `env` must be empty. + """ network: NotRequired[CreateEnvironmentRequestNetworkUnionParam] r"""Network configuration for the environment.""" sources: NotRequired[List[interactions_source.SourceParam]] @@ -63,6 +68,12 @@ class CreateEnvironmentRequestParam(TypedDict): class CreateEnvironmentRequest(BaseModel): r"""Request for `CreateEnvironment`.""" + from_environment: Optional[str] = None + r"""Optional. The source environment to copy/fork from. + Format: `environments/{environment_id}` or `{environment_id}`. + When specified, `sources` and `env` must be empty. + """ + network: Optional[CreateEnvironmentRequestNetworkUnion] = None r"""Network configuration for the environment.""" @@ -71,7 +82,7 @@ class CreateEnvironmentRequest(BaseModel): @model_serializer(mode="wrap") def serialize_model(self, handler): - optional_fields = set(["network", "sources"]) + optional_fields = set(["from_environment", "network", "sources"]) serialized = handler(self) m = {} diff --git a/google/genai/tests/gaos/test_environments_lifecycle.py b/google/genai/tests/gaos/test_environments_lifecycle.py index 224602593..39cb3c90c 100644 --- a/google/genai/tests/gaos/test_environments_lifecycle.py +++ b/google/genai/tests/gaos/test_environments_lifecycle.py @@ -19,6 +19,7 @@ from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer import json import threading +from typing import Any import pytest @@ -28,6 +29,9 @@ GeminiNextGenEnvironmentFiles, ) from ..._gaos.models.getenvironmentfiles import GetEnvironmentFilesRequest +from ..._gaos.types.environments.createenvironmentrequest import ( + CreateEnvironmentRequest, +) from ..._gaos.types.environments.environmentfile import EnvironmentFile from ..._gaos.types.environments.getenvironmentfilesresponse import ( GetEnvironmentFilesResponse, @@ -121,13 +125,18 @@ def test_python_environments_lifecycle_routes_through_google_genai_client( } ] ) + forked_environment = client.environments.create( + from_environment="environments/env_abc_1234", + ) client.environments.list() fetched = client.environments.get(id="env_abc_1234") client.environments.delete(id="env_abc_1234") assert environment.id == "env_abc_1234" + assert forked_environment.id == "env_abc_1234" assert fetched.id == "env_abc_1234" assert captured == [ + "POST /v1beta/environments", "POST /v1beta/environments", "GET /v1beta/environments", "GET /v1beta/environments/env_abc_1234", @@ -136,6 +145,49 @@ def test_python_environments_lifecycle_routes_through_google_genai_client( create_body = captured_bodies[0] assert create_body["sources"][0]["content"] == "print('hello')" + forked_body = captured_bodies[1] + assert forked_body["from_environment"] == "environments/env_abc_1234" + + finally: + server.shutdown() + thread.join() + server.server_close() + + +@pytest.mark.asyncio +async def test_python_environments_async_create_with_from_environment( + monkeypatch, +): + """Tests creating an environment asynchronously with from_environment.""" + monkeypatch.delenv("GOOGLE_GENAI_USE_VERTEXAI", raising=False) + captured: list[str] = [] + captured_bodies: list[dict[str, Any]] = [] + handler = type( + "Handler", + (_RecordingHandler,), + { + "captured": captured, + "captured_bodies": captured_bodies, + }, + ) + 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}", + }, + ) + + forked_environment = await client.aio.environments.create( + from_environment="environments/env_abc_1234", + ) + assert forked_environment.id == "env_abc_1234" + assert captured == ["POST /v1beta/environments"] + assert captured_bodies[0]["from_environment"] == "environments/env_abc_1234" finally: server.shutdown() @@ -372,6 +424,7 @@ def test_python_environments_types_and_models(): assert req.recursive is True assert req.api_version == "v1beta" - - - + create_req = CreateEnvironmentRequest( + from_environment="environments/env_abc_1234", + ) + assert create_req.from_environment == "environments/env_abc_1234"