fix: make project-id optional in all VertexAI generators (#1147)

* Make project-id optional in vertexai gens
deepset-ai · Oct 23, 2024 · 5b3b6c1 · 5b3b6c1
1 parent c09812c
commit 5b3b6c1
Show file tree

Hide file tree

Showing 10 changed files with 53 additions and 52 deletions.
diff --git a/.../google_vertex/src/haystack_integrations/components/generators/google_vertex/captioner.py b/.../google_vertex/src/haystack_integrations/components/generators/google_vertex/captioner.py
@@ -25,7 +25,7 @@ class VertexAIImageCaptioner:
     from haystack.dataclasses.byte_stream import ByteStream
     from haystack_integrations.components.generators.google_vertex import VertexAIImageCaptioner
 
-    captioner = VertexAIImageCaptioner(project_id=project_id)
+    captioner = VertexAIImageCaptioner()
 
     image = ByteStream(
         data=requests.get(
@@ -41,14 +41,16 @@ class VertexAIImageCaptioner:
     ```
     """
 
-    def __init__(self, *, model: str = "imagetext", project_id: str, location: Optional[str] = None, **kwargs):
+    def __init__(
+        self, *, model: str = "imagetext", project_id: Optional[str] = None, location: Optional[str] = None, **kwargs
+    ):
         """
         Generate image captions using a Google Vertex AI model.
 
         Authenticates using Google Cloud Application Default Credentials (ADCs).
         For more information see the official [Google documentation](https://cloud.google.com/docs/authentication/provide-credentials-adc).
 
-        :param project_id: ID of the GCP project to use.
+        :param project_id: ID of the GCP project to use. By default, it is set during Google Cloud authentication.
         :param model: Name of the model to use.
         :param location: The default location to use when making API calls, if not set uses us-central-1.
             Defaults to None.

diff --git a/...le_vertex/src/haystack_integrations/components/generators/google_vertex/code_generator.py b/...le_vertex/src/haystack_integrations/components/generators/google_vertex/code_generator.py
@@ -20,7 +20,7 @@ class VertexAICodeGenerator:
     ```python
         from haystack_integrations.components.generators.google_vertex import VertexAICodeGenerator
 
-        generator = VertexAICodeGenerator(project_id=project_id)
+        generator = VertexAICodeGenerator()
 
         result = generator.run(prefix="def to_json(data):")
 
@@ -45,14 +45,16 @@ class VertexAICodeGenerator:
     ```
     """
 
-    def __init__(self, *, model: str = "code-bison", project_id: str, location: Optional[str] = None, **kwargs):
+    def __init__(
+        self, *, model: str = "code-bison", project_id: Optional[str] = None, location: Optional[str] = None, **kwargs
+    ):
         """
         Generate code using a Google Vertex AI model.
 
         Authenticates using Google Cloud Application Default Credentials (ADCs).
         For more information see the official [Google documentation](https://cloud.google.com/docs/authentication/provide-credentials-adc).
 
-        :param project_id: ID of the GCP project to use.
+        :param project_id: ID of the GCP project to use. By default, it is set during Google Cloud authentication.
         :param model: Name of the model to use.
         :param location: The default location to use when making API calls, if not set uses us-central-1.
         :param kwargs: Additional keyword arguments to pass to the model.

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
@@ -24,20 +24,27 @@ class VertexAIImageGenerator:
 
     from haystack_integrations.components.generators.google_vertex import VertexAIImageGenerator
 
-    generator = VertexAIImageGenerator(project_id=project_id)
+    generator = VertexAIImageGenerator()
     result = generator.run(prompt="Generate an image of a cute cat")
     result["images"][0].to_file(Path("my_image.png"))
     ```
     """
 
-    def __init__(self, *, model: str = "imagegeneration", project_id: str, location: Optional[str] = None, **kwargs):
+    def __init__(
+        self,
+        *,
+        model: str = "imagegeneration",
+        project_id: Optional[str] = None,
+        location: Optional[str] = None,
+        **kwargs,
+    ):
         """
         Generates images using a Google Vertex AI model.
 
         Authenticates using Google Cloud Application Default Credentials (ADCs).
         For more information see the official [Google documentation](https://cloud.google.com/docs/authentication/provide-credentials-adc).
 
-        :param project_id: ID of the GCP project to use.
+        :param project_id: ID of the GCP project to use. By default, it is set during Google Cloud authentication.
         :param model: Name of the model to use.
         :param location: The default location to use when making API calls, if not set uses us-central-1.
         :param kwargs: Additional keyword arguments to pass to the model.

diff --git a/...ertex/src/haystack_integrations/components/generators/google_vertex/question_answering.py b/...ertex/src/haystack_integrations/components/generators/google_vertex/question_answering.py
@@ -23,7 +23,7 @@ class VertexAIImageQA:
     from haystack.dataclasses.byte_stream import ByteStream
     from haystack_integrations.components.generators.google_vertex import VertexAIImageQA
 
-    qa = VertexAIImageQA(project_id=project_id)
+    qa = VertexAIImageQA()
 
     image = ByteStream.from_file_path("dog.jpg")
 
@@ -35,14 +35,16 @@ class VertexAIImageQA:
     ```
     """
 
-    def __init__(self, *, model: str = "imagetext", project_id: str, location: Optional[str] = None, **kwargs):
+    def __init__(
+        self, *, model: str = "imagetext", project_id: Optional[str] = None, location: Optional[str] = None, **kwargs
+    ):
         """
         Answers questions about an image using a Google Vertex AI model.
 
         Authenticates using Google Cloud Application Default Credentials (ADCs).
         For more information see the official [Google documentation](https://cloud.google.com/docs/authentication/provide-credentials-adc).
 
-        :param project_id: ID of the GCP project to use.
+        :param project_id: ID of the GCP project to use. By default, it is set during Google Cloud authentication.
         :param model: Name of the model to use.
         :param location: The default location to use when making API calls, if not set uses us-central-1.
         :param kwargs: Additional keyword arguments to pass to the model.

diff --git a/...le_vertex/src/haystack_integrations/components/generators/google_vertex/text_generator.py b/...le_vertex/src/haystack_integrations/components/generators/google_vertex/text_generator.py
@@ -25,7 +25,7 @@ class VertexAITextGenerator:
     ```python
         from haystack_integrations.components.generators.google_vertex import VertexAITextGenerator
 
-        generator = VertexAITextGenerator(project_id=project_id)
+        generator = VertexAITextGenerator()
         res = generator.run("Tell me a good interview question for a software engineer.")
 
         print(res["replies"][0])
@@ -45,14 +45,16 @@ class VertexAITextGenerator:
     ```
     """
 
-    def __init__(self, *, model: str = "text-bison", project_id: str, location: Optional[str] = None, **kwargs):
+    def __init__(
+        self, *, model: str = "text-bison", project_id: Optional[str] = None, location: Optional[str] = None, **kwargs
+    ):
         """
         Generate text using a Google Vertex AI model.
 
         Authenticates using Google Cloud Application Default Credentials (ADCs).
         For more information see the official [Google documentation](https://cloud.google.com/docs/authentication/provide-credentials-adc).
 
-        :param project_id: ID of the GCP project to use.
+        :param project_id: ID of the GCP project to use. By default, it is set during Google Cloud authentication.
         :param model: Name of the model to use.
         :param location: The default location to use when making API calls, if not set uses us-central-1.
         :param kwargs: Additional keyword arguments to pass to the model.

diff --git a/integrations/google_vertex/tests/test_captioner.py b/integrations/google_vertex/tests/test_captioner.py
@@ -22,14 +22,12 @@ def test_init(mock_model_class, mock_vertexai):
 @patch("haystack_integrations.components.generators.google_vertex.captioner.vertexai")
 @patch("haystack_integrations.components.generators.google_vertex.captioner.ImageTextModel")
 def test_to_dict(_mock_model_class, _mock_vertexai):
-    captioner = VertexAIImageCaptioner(
-        model="imagetext", project_id="myproject-123456", number_of_results=1, language="it"
-    )
+    captioner = VertexAIImageCaptioner(model="imagetext", number_of_results=1, language="it")
     assert captioner.to_dict() == {
         "type": "haystack_integrations.components.generators.google_vertex.captioner.VertexAIImageCaptioner",
         "init_parameters": {
             "model": "imagetext",
-            "project_id": "myproject-123456",
+            "project_id": None,
             "location": None,
             "number_of_results": 1,
             "language": "it",
@@ -45,14 +43,15 @@ def test_from_dict(_mock_model_class, _mock_vertexai):
             "type": "haystack_integrations.components.generators.google_vertex.captioner.VertexAIImageCaptioner",
             "init_parameters": {
                 "model": "imagetext",
-                "project_id": "myproject-123456",
+                "project_id": None,
+                "location": None,
                 "number_of_results": 1,
                 "language": "it",
             },
         }
     )
     assert captioner._model_name == "imagetext"
-    assert captioner._project_id == "myproject-123456"
+    assert captioner._project_id is None
     assert captioner._location is None
     assert captioner._kwargs == {"number_of_results": 1, "language": "it"}
     assert captioner._model is not None
@@ -63,9 +62,7 @@ def test_from_dict(_mock_model_class, _mock_vertexai):
 def test_run_calls_get_captions(mock_model_class, _mock_vertexai):
     mock_model = Mock()
     mock_model_class.from_pretrained.return_value = mock_model
-    captioner = VertexAIImageCaptioner(
-        model="imagetext", project_id="myproject-123456", number_of_results=1, language="it"
-    )
+    captioner = VertexAIImageCaptioner(model="imagetext", number_of_results=1, language="it")
 
     image = ByteStream(data=b"image data")
     captioner.run(image=image)

diff --git a/integrations/google_vertex/tests/test_code_generator.py b/integrations/google_vertex/tests/test_code_generator.py
@@ -22,14 +22,12 @@ def test_init(mock_model_class, mock_vertexai):
 @patch("haystack_integrations.components.generators.google_vertex.code_generator.vertexai")
 @patch("haystack_integrations.components.generators.google_vertex.code_generator.CodeGenerationModel")
 def test_to_dict(_mock_model_class, _mock_vertexai):
-    generator = VertexAICodeGenerator(
-        model="code-bison", project_id="myproject-123456", candidate_count=3, temperature=0.5
-    )
+    generator = VertexAICodeGenerator(model="code-bison", candidate_count=3, temperature=0.5)
     assert generator.to_dict() == {
         "type": "haystack_integrations.components.generators.google_vertex.code_generator.VertexAICodeGenerator",
         "init_parameters": {
             "model": "code-bison",
-            "project_id": "myproject-123456",
+            "project_id": None,
             "location": None,
             "candidate_count": 3,
             "temperature": 0.5,
@@ -45,14 +43,15 @@ def test_from_dict(_mock_model_class, _mock_vertexai):
             "type": "haystack_integrations.components.generators.google_vertex.code_generator.VertexAICodeGenerator",
             "init_parameters": {
                 "model": "code-bison",
-                "project_id": "myproject-123456",
+                "project_id": None,
+                "location": None,
                 "candidate_count": 2,
                 "temperature": 0.5,
             },
         }
     )
     assert generator._model_name == "code-bison"
-    assert generator._project_id == "myproject-123456"
+    assert generator._project_id is None
     assert generator._location is None
     assert generator._kwargs == {"candidate_count": 2, "temperature": 0.5}
     assert generator._model is not None
@@ -64,9 +63,7 @@ def test_run_calls_predict(mock_model_class, _mock_vertexai):
     mock_model = Mock()
     mock_model.predict.return_value = TextGenerationResponse("answer", None)
     mock_model_class.from_pretrained.return_value = mock_model
-    generator = VertexAICodeGenerator(
-        model="code-bison", project_id="myproject-123456", candidate_count=1, temperature=0.5
-    )
+    generator = VertexAICodeGenerator(model="code-bison", candidate_count=1, temperature=0.5)
 
     prefix = "def print_json(data):\n"
     generator.run(prefix=prefix)

diff --git a/integrations/google_vertex/tests/test_image_generator.py b/integrations/google_vertex/tests/test_image_generator.py
@@ -30,15 +30,14 @@ def test_init(mock_model_class, mock_vertexai):
 def test_to_dict(_mock_model_class, _mock_vertexai):
     generator = VertexAIImageGenerator(
         model="imagetext",
-        project_id="myproject-123456",
         guidance_scale=12,
         number_of_images=3,
     )
     assert generator.to_dict() == {
         "type": "haystack_integrations.components.generators.google_vertex.image_generator.VertexAIImageGenerator",
         "init_parameters": {
             "model": "imagetext",
-            "project_id": "myproject-123456",
+            "project_id": None,
             "location": None,
             "guidance_scale": 12,
             "number_of_images": 3,
@@ -54,15 +53,15 @@ def test_from_dict(_mock_model_class, _mock_vertexai):
             "type": "haystack_integrations.components.generators.google_vertex.image_generator.VertexAIImageGenerator",
             "init_parameters": {
                 "model": "imagetext",
-                "project_id": "myproject-123456",
+                "project_id": None,
                 "location": None,
                 "guidance_scale": 12,
                 "number_of_images": 3,
             },
         }
     )
     assert generator._model_name == "imagetext"
-    assert generator._project_id == "myproject-123456"
+    assert generator._project_id is None
     assert generator._location is None
     assert generator._kwargs == {
         "guidance_scale": 12,
@@ -78,7 +77,6 @@ def test_run_calls_generate_images(mock_model_class, _mock_vertexai):
     mock_model_class.from_pretrained.return_value = mock_model
     generator = VertexAIImageGenerator(
         model="imagetext",
-        project_id="myproject-123456",
         guidance_scale=12,
         number_of_images=3,
     )

diff --git a/integrations/google_vertex/tests/test_question_answering.py b/integrations/google_vertex/tests/test_question_answering.py
@@ -26,14 +26,13 @@ def test_init(mock_model_class, mock_vertexai):
 def test_to_dict(_mock_model_class, _mock_vertexai):
     generator = VertexAIImageQA(
         model="imagetext",
-        project_id="myproject-123456",
         number_of_results=3,
     )
     assert generator.to_dict() == {
         "type": "haystack_integrations.components.generators.google_vertex.question_answering.VertexAIImageQA",
         "init_parameters": {
             "model": "imagetext",
-            "project_id": "myproject-123456",
+            "project_id": None,
             "location": None,
             "number_of_results": 3,
         },
@@ -48,14 +47,14 @@ def test_from_dict(_mock_model_class, _mock_vertexai):
             "type": "haystack_integrations.components.generators.google_vertex.question_answering.VertexAIImageQA",
             "init_parameters": {
                 "model": "imagetext",
-                "project_id": "myproject-123456",
+                "project_id": None,
                 "location": None,
                 "number_of_results": 3,
             },
         }
     )
     assert generator._model_name == "imagetext"
-    assert generator._project_id == "myproject-123456"
+    assert generator._project_id is None
     assert generator._location is None
     assert generator._kwargs == {"number_of_results": 3}
 
@@ -68,7 +67,6 @@ def test_run_calls_ask_question(mock_model_class, _mock_vertexai):
     mock_model_class.from_pretrained.return_value = mock_model
     generator = VertexAIImageQA(
         model="imagetext",
-        project_id="myproject-123456",
         number_of_results=3,
     )
 

diff --git a/integrations/google_vertex/tests/test_text_generator.py b/integrations/google_vertex/tests/test_text_generator.py
@@ -24,14 +24,12 @@ def test_init(mock_model_class, mock_vertexai):
 @patch("haystack_integrations.components.generators.google_vertex.text_generator.TextGenerationModel")
 def test_to_dict(_mock_model_class, _mock_vertexai):
     grounding_source = GroundingSource.VertexAISearch("1234", "us-central-1")
-    generator = VertexAITextGenerator(
-        model="text-bison", project_id="myproject-123456", temperature=0.2, grounding_source=grounding_source
-    )
+    generator = VertexAITextGenerator(model="text-bison", temperature=0.2, grounding_source=grounding_source)
     assert generator.to_dict() == {
         "type": "haystack_integrations.components.generators.google_vertex.text_generator.VertexAITextGenerator",
         "init_parameters": {
             "model": "text-bison",
-            "project_id": "myproject-123456",
+            "project_id": None,
             "location": None,
             "temperature": 0.2,
             "grounding_source": {
@@ -55,7 +53,7 @@ def test_from_dict(_mock_model_class, _mock_vertexai):
             "type": "haystack_integrations.components.generators.google_vertex.text_generator.VertexAITextGenerator",
             "init_parameters": {
                 "model": "text-bison",
-                "project_id": "myproject-123456",
+                "project_id": None,
                 "location": None,
                 "temperature": 0.2,
                 "grounding_source": {
@@ -71,7 +69,7 @@ def test_from_dict(_mock_model_class, _mock_vertexai):
         }
     )
     assert generator._model_name == "text-bison"
-    assert generator._project_id == "myproject-123456"
+    assert generator._project_id is None
     assert generator._location is None
     assert generator._kwargs == {
         "temperature": 0.2,
@@ -86,9 +84,7 @@ def test_run_calls_get_captions(mock_model_class, _mock_vertexai):
     mock_model.predict.return_value = MagicMock()
     mock_model_class.from_pretrained.return_value = mock_model
     grounding_source = GroundingSource.VertexAISearch("1234", "us-central-1")
-    generator = VertexAITextGenerator(
-        model="text-bison", project_id="myproject-123456", temperature=0.2, grounding_source=grounding_source
-    )
+    generator = VertexAITextGenerator(model="text-bison", temperature=0.2, grounding_source=grounding_source)
 
     prompt = "What is the answer?"
     generator.run(prompt=prompt)