@@ -23,7 +23,7 @@
 from google import auth
 from google.cloud import aiplatform
 from tests.system.aiplatform import e2e_base
-from vertexai.preview import generative_models
+from vertexai import generative_models
 
 
 class TestGenerativeModels(e2e_base.TestEndToEnd):

@@ -22,7 +22,8 @@
 
 import vertexai
 from google.cloud.aiplatform import initializer
-from vertexai.preview import generative_models
+from vertexai import generative_models
+from vertexai.preview import generative_models as preview_generative_models
 from vertexai.generative_models._generative_models import (
 prediction_service,
 gapic_prediction_service_types,
@@ -231,7 +232,11 @@ def teardown_method(self):
 attribute="generate_content",
 new=mock_generate_content,
 )
-def test_generate_content(self):
+@pytest.mark.parametrize(
+"generative_models",
+[generative_models, preview_generative_models],
+)
+def test_generate_content(self, generative_models: generative_models):
 model = generative_models.GenerativeModel("gemini-pro")
 response = model.generate_content("Why is sky blue?")
 assert response.text
@@ -254,7 +259,11 @@ def test_generate_content(self):
 attribute="stream_generate_content",
 new=mock_stream_generate_content,
 )
-def test_generate_content_(self):
+@pytest.mark.parametrize(
+"generative_models",
+[generative_models, preview_generative_models],
+)
+def test_generate_content_(self, generative_models: generative_models):
 model = generative_models.GenerativeModel("gemini-pro")
 stream = model.generate_content("Why is sky blue?", stream=True)
 for chunk in stream:
@@ -265,7 +274,11 @@ def test_generate_content_(self):
 attribute="generate_content",
 new=mock_generate_content,
 )
-def test_chat_send_message(self):
+@pytest.mark.parametrize(
+"generative_models",
+[generative_models, preview_generative_models],
+)
+def test_chat_send_message(self, generative_models: generative_models):
 model = generative_models.GenerativeModel("gemini-pro")
 chat = model.start_chat()
 response1 = chat.send_message("Why is sky blue?")
@@ -278,7 +291,11 @@ def test_chat_send_message(self):
 attribute="generate_content",
 new=mock_generate_content,
 )
-def test_chat_function_calling(self):
+@pytest.mark.parametrize(
+"generative_models",
+[generative_models, preview_generative_models],
+)
+def test_chat_function_calling(self, generative_models: generative_models):
 get_current_weather_func = generative_models.FunctionDeclaration(
 name="get_current_weather",
 description="Get the current weather in a given location",

@@ -0,0 +1,51 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+"""Classes for working with the Gemini models."""
+
+# We just want to re-export certain classes
+# pylint: disable=g-multiple-import,g-importing-member
+from vertexai.generative_models._generative_models import (
+GenerativeModel,
+GenerationConfig,
+GenerationResponse,
+Candidate,
+ChatSession,
+Content,
+FinishReason,
+FunctionDeclaration,
+HarmCategory,
+HarmBlockThreshold,
+Image,
+Part,
+ResponseBlockedError,
+Tool,
+)
+
+__all__ = [
+"GenerationConfig",
+"GenerativeModel",
+"GenerationResponse",
+"Candidate",
+"ChatSession",
+"Content",
+"FinishReason",
+"FunctionDeclaration",
+"HarmCategory",
+"HarmBlockThreshold",
+"Image",
+"Part",
+"ResponseBlockedError",
+"Tool",
+]

@@ -1896,6 +1896,10 @@ def _tune_model(
 return job
 
 
+class GenerativeModel(_GenerativeModel):
+__module__ = "vertexai.generative_models"
+
+
 class _PreviewGenerativeModel(_GenerativeModel, _TunableGenerativeModelMixin):
 __name__ = "GenerativeModel"
 __module__ = "vertexai.preview.generative_models"