Update docstrings to remove vertex.preview (#1074)

deepset-ai · Sep 11, 2024 · e7de37e · e7de37e
1 parent 6daefc6
commit e7de37e
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 10 deletions.
diff --git a/...oogle_vertex/src/haystack_integrations/components/generators/google_vertex/chat/gemini.py b/...oogle_vertex/src/haystack_integrations/components/generators/google_vertex/chat/gemini.py
@@ -67,14 +67,14 @@ def __init__(
         :param location: The default location to use when making API calls, if not set uses us-central-1.
             Defaults to None.
         :param generation_config: Configuration for the generation process.
-            See the [GenerationConfig documentation](https://cloud.google.com/python/docs/reference/aiplatform/latest/vertexai.preview.generative_models.GenerationConfig
+            See the [GenerationConfig documentation](https://cloud.google.com/python/docs/reference/aiplatform/latest/vertexai.generative_models.GenerationConfig
             for a list of supported arguments.
         :param safety_settings: Safety settings to use when generating content. See the documentation
-            for [HarmBlockThreshold](https://cloud.google.com/python/docs/reference/aiplatform/latest/vertexai.preview.generative_models.HarmBlockThreshold)
-            and [HarmCategory](https://cloud.google.com/python/docs/reference/aiplatform/latest/vertexai.preview.generative_models.HarmCategory)
+            for [HarmBlockThreshold](https://cloud.google.com/python/docs/reference/aiplatform/latest/vertexai.generative_models.HarmBlockThreshold)
+            and [HarmCategory](https://cloud.google.com/python/docs/reference/aiplatform/latest/vertexai.generative_models.HarmCategory)
             for more details.
         :param tools: List of tools to use when generating content. See the documentation for
-            [Tool](https://cloud.google.com/python/docs/reference/aiplatform/latest/vertexai.preview.generative_models.Tool)
+            [Tool](https://cloud.google.com/python/docs/reference/aiplatform/latest/vertexai.generative_models.Tool)
             the list of supported arguments.
         :param streaming_callback: A callback function that is called when a new token is received from
             the  stream. The callback function accepts StreamingChunk as an argument.

diff --git a/...ons/google_vertex/src/haystack_integrations/components/generators/google_vertex/gemini.py b/...ons/google_vertex/src/haystack_integrations/components/generators/google_vertex/gemini.py
@@ -70,7 +70,7 @@ def __init__(
         :param model: Name of the model to use. For available models, see https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models.
         :param location: The default location to use when making API calls, if not set uses us-central-1.
         :param generation_config: The generation config to use.
-            Can either be a [`GenerationConfig`](https://cloud.google.com/python/docs/reference/aiplatform/latest/vertexai.preview.generative_models.GenerationConfig)
+            Can either be a [`GenerationConfig`](https://cloud.google.com/python/docs/reference/aiplatform/latest/vertexai.generative_models.GenerationConfig)
             object or a dictionary of parameters.
             Accepted fields are:
                 - temperature
@@ -80,11 +80,11 @@ def __init__(
                 - max_output_tokens
                 - stop_sequences
         :param safety_settings: The safety settings to use. See the documentation
-            for [HarmBlockThreshold](https://cloud.google.com/python/docs/reference/aiplatform/latest/vertexai.preview.generative_models.HarmBlockThreshold)
-            and [HarmCategory](https://cloud.google.com/python/docs/reference/aiplatform/latest/vertexai.preview.generative_models.HarmCategory)
+            for [HarmBlockThreshold](https://cloud.google.com/python/docs/reference/aiplatform/latest/vertexai.generative_models.HarmBlockThreshold)
+            and [HarmCategory](https://cloud.google.com/python/docs/reference/aiplatform/latest/vertexai.generative_models.HarmCategory)
             for more details.
         :param tools: List of tools to use when generating content. See the documentation for
-            [Tool](https://cloud.google.com/python/docs/reference/aiplatform/latest/vertexai.preview.generative_models.Tool)
+            [Tool](https://cloud.google.com/python/docs/reference/aiplatform/latest/vertexai.generative_models.Tool)
             the list of supported arguments.
         :param streaming_callback: A callback function that is called when a new token is received from the stream.
             The callback function accepts StreamingChunk as an argument.

diff --git a/...e_vertex/src/haystack_integrations/components/generators/google_vertex/image_generator.py b/...e_vertex/src/haystack_integrations/components/generators/google_vertex/image_generator.py
@@ -5,7 +5,7 @@
 from haystack.core.component import component
 from haystack.core.serialization import default_from_dict, default_to_dict
 from haystack.dataclasses.byte_stream import ByteStream
-from vertexai.preview.vision_models import ImageGenerationModel
+from vertexai.vision_models import ImageGenerationModel
 
 logger = logging.getLogger(__name__)
 

diff --git a/integrations/google_vertex/tests/test_image_generator.py b/integrations/google_vertex/tests/test_image_generator.py
@@ -1,6 +1,6 @@
 from unittest.mock import Mock, patch
 
-from vertexai.preview.vision_models import ImageGenerationResponse
+from vertexai.vision_models import ImageGenerationResponse
 
 from haystack_integrations.components.generators.google_vertex import VertexAIImageGenerator