From 4118b629b2053a0d09b872343fd17a1ca90ab677 Mon Sep 17 00:00:00 2001 From: Massimiliano Pippi Date: Fri, 26 Jan 2024 16:02:02 +0100 Subject: [PATCH] change import paths (#273) --- integrations/google_vertex/pyproject.toml | 19 ++++++++-------- .../src/google_vertex_haystack/__init__.py | 3 --- .../generators/__init__.py | 3 --- .../generators/google_vertex/__init__.py | 20 +++++++++++++++++ .../generators/google_vertex}/captioner.py | 0 .../generators/google_vertex/chat/__init__.py | 0 .../generators/google_vertex}/chat/gemini.py | 0 .../google_vertex}/code_generator.py | 0 .../generators/google_vertex}/gemini.py | 0 .../google_vertex}/image_generator.py | 0 .../google_vertex}/question_answering.py | 0 .../google_vertex}/text_generator.py | 0 .../google_vertex/tests/test_captioner.py | 22 +++++++++---------- .../tests/test_code_generator.py | 22 +++++++++---------- .../tests/test_image_generator.py | 22 +++++++++---------- .../tests/test_question_answering.py | 22 +++++++++---------- .../tests/test_text_generator.py | 22 +++++++++---------- 17 files changed, 85 insertions(+), 70 deletions(-) delete mode 100644 integrations/google_vertex/src/google_vertex_haystack/__init__.py delete mode 100644 integrations/google_vertex/src/google_vertex_haystack/generators/__init__.py create mode 100644 integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/__init__.py rename integrations/google_vertex/src/{google_vertex_haystack/generators => haystack_integrations/components/generators/google_vertex}/captioner.py (100%) create mode 100644 integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/chat/__init__.py rename integrations/google_vertex/src/{google_vertex_haystack/generators => haystack_integrations/components/generators/google_vertex}/chat/gemini.py (100%) rename integrations/google_vertex/src/{google_vertex_haystack/generators => haystack_integrations/components/generators/google_vertex}/code_generator.py (100%) rename integrations/google_vertex/src/{google_vertex_haystack/generators => haystack_integrations/components/generators/google_vertex}/gemini.py (100%) rename integrations/google_vertex/src/{google_vertex_haystack/generators => haystack_integrations/components/generators/google_vertex}/image_generator.py (100%) rename integrations/google_vertex/src/{google_vertex_haystack/generators => haystack_integrations/components/generators/google_vertex}/question_answering.py (100%) rename integrations/google_vertex/src/{google_vertex_haystack/generators => haystack_integrations/components/generators/google_vertex}/text_generator.py (100%) diff --git a/integrations/google_vertex/pyproject.toml b/integrations/google_vertex/pyproject.toml index 1d15a4270..ecd509f15 100644 --- a/integrations/google_vertex/pyproject.toml +++ b/integrations/google_vertex/pyproject.toml @@ -33,6 +33,9 @@ Documentation = "https://github.com/deepset-ai/haystack-core-integrations/tree/m Issues = "https://github.com/deepset-ai/haystack-core-integrations/issues" Source = "https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/google_vertex" +[tool.hatch.build.targets.wheel] +packages = ["src/haystack_integrations"] + [tool.hatch.version] source = "vcs" tag-pattern = 'integrations\/google_vertex-v(?P.*)' @@ -69,7 +72,7 @@ dependencies = [ "ruff>=0.0.243", ] [tool.hatch.envs.lint.scripts] -typing = "mypy --install-types --non-interactive {args:src/google_vertex_haystack tests}" +typing = "mypy --install-types --non-interactive --explicit-package-bases {args:src/ tests}" style = [ "ruff {args:.}", "black --check --diff {args:.}", @@ -132,26 +135,23 @@ unfixable = [ ] [tool.ruff.isort] -known-first-party = ["google_vertex_haystack"] +known-first-party = ["haystack_integrations"] [tool.ruff.flake8-tidy-imports] -ban-relative-imports = "all" +ban-relative-imports = "parents" [tool.ruff.per-file-ignores] # Tests can use magic values, assertions, and relative imports "tests/**/*" = ["PLR2004", "S101", "TID252"] [tool.coverage.run] -source_pkgs = ["google_vertex_haystack", "tests"] +source_pkgs = ["haystack_integrations", "tests"] branch = true parallel = true -omit = [ - "src/google_vertex_haystack/__about__.py", -] [tool.coverage.paths] -google_vertex_haystack = ["src/google_vertex_haystack", "*/google_vertex/src/google_vertex_haystack"] -tests = ["tests", "*/google_vertex_haystack/tests"] +google_vertex_haystack = ["src/"] +tests = ["tests"] [tool.coverage.report] exclude_lines = [ @@ -164,6 +164,7 @@ exclude_lines = [ module = [ "vertexai.*", "haystack.*", + "haystack_integrations.*", "pytest.*", "numpy.*", ] diff --git a/integrations/google_vertex/src/google_vertex_haystack/__init__.py b/integrations/google_vertex/src/google_vertex_haystack/__init__.py deleted file mode 100644 index e873bc332..000000000 --- a/integrations/google_vertex/src/google_vertex_haystack/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -# SPDX-FileCopyrightText: 2023-present deepset GmbH -# -# SPDX-License-Identifier: Apache-2.0 diff --git a/integrations/google_vertex/src/google_vertex_haystack/generators/__init__.py b/integrations/google_vertex/src/google_vertex_haystack/generators/__init__.py deleted file mode 100644 index e873bc332..000000000 --- a/integrations/google_vertex/src/google_vertex_haystack/generators/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -# SPDX-FileCopyrightText: 2023-present deepset GmbH -# -# SPDX-License-Identifier: Apache-2.0 diff --git a/integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/__init__.py b/integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/__init__.py new file mode 100644 index 000000000..07c2a5260 --- /dev/null +++ b/integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/__init__.py @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: 2023-present deepset GmbH +# +# SPDX-License-Identifier: Apache-2.0 +from .captioner import VertexAIImageCaptioner +from .chat.gemini import VertexAIGeminiChatGenerator +from .code_generator import VertexAICodeGenerator +from .gemini import VertexAIGeminiGenerator +from .image_generator import VertexAIImageGenerator +from .question_answering import VertexAIImageQA +from .text_generator import VertexAITextGenerator + +__all__ = [ + "VertexAICodeGenerator", + "VertexAIGeminiGenerator", + "VertexAIGeminiChatGenerator", + "VertexAIImageCaptioner", + "VertexAIImageGenerator", + "VertexAIImageQA", + "VertexAITextGenerator", +] diff --git a/integrations/google_vertex/src/google_vertex_haystack/generators/captioner.py b/integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/captioner.py similarity index 100% rename from integrations/google_vertex/src/google_vertex_haystack/generators/captioner.py rename to integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/captioner.py diff --git a/integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/chat/__init__.py b/integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/chat/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/integrations/google_vertex/src/google_vertex_haystack/generators/chat/gemini.py b/integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/chat/gemini.py similarity index 100% rename from integrations/google_vertex/src/google_vertex_haystack/generators/chat/gemini.py rename to integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/chat/gemini.py diff --git a/integrations/google_vertex/src/google_vertex_haystack/generators/code_generator.py b/integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/code_generator.py similarity index 100% rename from integrations/google_vertex/src/google_vertex_haystack/generators/code_generator.py rename to integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/code_generator.py diff --git a/integrations/google_vertex/src/google_vertex_haystack/generators/gemini.py b/integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/gemini.py similarity index 100% rename from integrations/google_vertex/src/google_vertex_haystack/generators/gemini.py rename to integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/gemini.py diff --git a/integrations/google_vertex/src/google_vertex_haystack/generators/image_generator.py b/integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/image_generator.py similarity index 100% rename from integrations/google_vertex/src/google_vertex_haystack/generators/image_generator.py rename to integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/image_generator.py diff --git a/integrations/google_vertex/src/google_vertex_haystack/generators/question_answering.py b/integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/question_answering.py similarity index 100% rename from integrations/google_vertex/src/google_vertex_haystack/generators/question_answering.py rename to integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/question_answering.py diff --git a/integrations/google_vertex/src/google_vertex_haystack/generators/text_generator.py b/integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/text_generator.py similarity index 100% rename from integrations/google_vertex/src/google_vertex_haystack/generators/text_generator.py rename to integrations/google_vertex/src/haystack_integrations/components/generators/google_vertex/text_generator.py diff --git a/integrations/google_vertex/tests/test_captioner.py b/integrations/google_vertex/tests/test_captioner.py index bc7e4f829..26249dbee 100644 --- a/integrations/google_vertex/tests/test_captioner.py +++ b/integrations/google_vertex/tests/test_captioner.py @@ -2,11 +2,11 @@ from haystack.dataclasses.byte_stream import ByteStream -from google_vertex_haystack.generators.captioner import VertexAIImageCaptioner +from haystack_integrations.components.generators.google_vertex import VertexAIImageCaptioner -@patch("google_vertex_haystack.generators.captioner.vertexai") -@patch("google_vertex_haystack.generators.captioner.ImageTextModel") +@patch("haystack_integrations.components.generators.google_vertex.captioner.vertexai") +@patch("haystack_integrations.components.generators.google_vertex.captioner.ImageTextModel") def test_init(mock_model_class, mock_vertexai): captioner = VertexAIImageCaptioner( model="imagetext", project_id="myproject-123456", number_of_results=1, language="it" @@ -19,14 +19,14 @@ def test_init(mock_model_class, mock_vertexai): assert captioner._kwargs == {"number_of_results": 1, "language": "it"} -@patch("google_vertex_haystack.generators.captioner.vertexai") -@patch("google_vertex_haystack.generators.captioner.ImageTextModel") +@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" ) assert captioner.to_dict() == { - "type": "google_vertex_haystack.generators.captioner.VertexAIImageCaptioner", + "type": "haystack_integrations.components.generators.google_vertex.captioner.VertexAIImageCaptioner", "init_parameters": { "model": "imagetext", "project_id": "myproject-123456", @@ -37,12 +37,12 @@ def test_to_dict(_mock_model_class, _mock_vertexai): } -@patch("google_vertex_haystack.generators.captioner.vertexai") -@patch("google_vertex_haystack.generators.captioner.ImageTextModel") +@patch("haystack_integrations.components.generators.google_vertex.captioner.vertexai") +@patch("haystack_integrations.components.generators.google_vertex.captioner.ImageTextModel") def test_from_dict(_mock_model_class, _mock_vertexai): captioner = VertexAIImageCaptioner.from_dict( { - "type": "google_vertex_haystack.generators.captioner.VertexAIImageCaptioner", + "type": "haystack_integrations.components.generators.google_vertex.captioner.VertexAIImageCaptioner", "init_parameters": { "model": "imagetext", "project_id": "myproject-123456", @@ -58,8 +58,8 @@ def test_from_dict(_mock_model_class, _mock_vertexai): assert captioner._model is not None -@patch("google_vertex_haystack.generators.captioner.vertexai") -@patch("google_vertex_haystack.generators.captioner.ImageTextModel") +@patch("haystack_integrations.components.generators.google_vertex.captioner.vertexai") +@patch("haystack_integrations.components.generators.google_vertex.captioner.ImageTextModel") def test_run_calls_get_captions(mock_model_class, _mock_vertexai): mock_model = Mock() mock_model_class.from_pretrained.return_value = mock_model diff --git a/integrations/google_vertex/tests/test_code_generator.py b/integrations/google_vertex/tests/test_code_generator.py index c2a2e5aa9..129954062 100644 --- a/integrations/google_vertex/tests/test_code_generator.py +++ b/integrations/google_vertex/tests/test_code_generator.py @@ -2,11 +2,11 @@ from vertexai.language_models import TextGenerationResponse -from google_vertex_haystack.generators.code_generator import VertexAICodeGenerator +from haystack_integrations.components.generators.google_vertex import VertexAICodeGenerator -@patch("google_vertex_haystack.generators.code_generator.vertexai") -@patch("google_vertex_haystack.generators.code_generator.CodeGenerationModel") +@patch("haystack_integrations.components.generators.google_vertex.code_generator.vertexai") +@patch("haystack_integrations.components.generators.google_vertex.code_generator.CodeGenerationModel") def test_init(mock_model_class, mock_vertexai): generator = VertexAICodeGenerator( model="code-bison", project_id="myproject-123456", candidate_count=3, temperature=0.5 @@ -19,14 +19,14 @@ def test_init(mock_model_class, mock_vertexai): assert generator._kwargs == {"candidate_count": 3, "temperature": 0.5} -@patch("google_vertex_haystack.generators.code_generator.vertexai") -@patch("google_vertex_haystack.generators.code_generator.CodeGenerationModel") +@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 ) assert generator.to_dict() == { - "type": "google_vertex_haystack.generators.code_generator.VertexAICodeGenerator", + "type": "haystack_integrations.components.generators.google_vertex.code_generator.VertexAICodeGenerator", "init_parameters": { "model": "code-bison", "project_id": "myproject-123456", @@ -37,12 +37,12 @@ def test_to_dict(_mock_model_class, _mock_vertexai): } -@patch("google_vertex_haystack.generators.code_generator.vertexai") -@patch("google_vertex_haystack.generators.code_generator.CodeGenerationModel") +@patch("haystack_integrations.components.generators.google_vertex.code_generator.vertexai") +@patch("haystack_integrations.components.generators.google_vertex.code_generator.CodeGenerationModel") def test_from_dict(_mock_model_class, _mock_vertexai): generator = VertexAICodeGenerator.from_dict( { - "type": "google_vertex_haystack.generators.code_generator.VertexAICodeGenerator", + "type": "haystack_integrations.components.generators.google_vertex.code_generator.VertexAICodeGenerator", "init_parameters": { "model": "code-bison", "project_id": "myproject-123456", @@ -58,8 +58,8 @@ def test_from_dict(_mock_model_class, _mock_vertexai): assert generator._model is not None -@patch("google_vertex_haystack.generators.code_generator.vertexai") -@patch("google_vertex_haystack.generators.code_generator.CodeGenerationModel") +@patch("haystack_integrations.components.generators.google_vertex.code_generator.vertexai") +@patch("haystack_integrations.components.generators.google_vertex.code_generator.CodeGenerationModel") def test_run_calls_predict(mock_model_class, _mock_vertexai): mock_model = Mock() mock_model.predict.return_value = TextGenerationResponse("answer", None) diff --git a/integrations/google_vertex/tests/test_image_generator.py b/integrations/google_vertex/tests/test_image_generator.py index 1c5381a48..42cc0a0a3 100644 --- a/integrations/google_vertex/tests/test_image_generator.py +++ b/integrations/google_vertex/tests/test_image_generator.py @@ -2,11 +2,11 @@ from vertexai.preview.vision_models import ImageGenerationResponse -from google_vertex_haystack.generators.image_generator import VertexAIImageGenerator +from haystack_integrations.components.generators.google_vertex import VertexAIImageGenerator -@patch("google_vertex_haystack.generators.image_generator.vertexai") -@patch("google_vertex_haystack.generators.image_generator.ImageGenerationModel") +@patch("haystack_integrations.components.generators.google_vertex.image_generator.vertexai") +@patch("haystack_integrations.components.generators.google_vertex.image_generator.ImageGenerationModel") def test_init(mock_model_class, mock_vertexai): generator = VertexAIImageGenerator( model="imagetext", @@ -25,8 +25,8 @@ def test_init(mock_model_class, mock_vertexai): } -@patch("google_vertex_haystack.generators.image_generator.vertexai") -@patch("google_vertex_haystack.generators.image_generator.ImageGenerationModel") +@patch("haystack_integrations.components.generators.google_vertex.image_generator.vertexai") +@patch("haystack_integrations.components.generators.google_vertex.image_generator.ImageGenerationModel") def test_to_dict(_mock_model_class, _mock_vertexai): generator = VertexAIImageGenerator( model="imagetext", @@ -35,7 +35,7 @@ def test_to_dict(_mock_model_class, _mock_vertexai): number_of_images=3, ) assert generator.to_dict() == { - "type": "google_vertex_haystack.generators.image_generator.VertexAIImageGenerator", + "type": "haystack_integrations.components.generators.google_vertex.image_generator.VertexAIImageGenerator", "init_parameters": { "model": "imagetext", "project_id": "myproject-123456", @@ -46,12 +46,12 @@ def test_to_dict(_mock_model_class, _mock_vertexai): } -@patch("google_vertex_haystack.generators.image_generator.vertexai") -@patch("google_vertex_haystack.generators.image_generator.ImageGenerationModel") +@patch("haystack_integrations.components.generators.google_vertex.image_generator.vertexai") +@patch("haystack_integrations.components.generators.google_vertex.image_generator.ImageGenerationModel") def test_from_dict(_mock_model_class, _mock_vertexai): generator = VertexAIImageGenerator.from_dict( { - "type": "google_vertex_haystack.generators.image_generator.VertexAIImageGenerator", + "type": "haystack_integrations.components.generators.google_vertex.image_generator.VertexAIImageGenerator", "init_parameters": { "model": "imagetext", "project_id": "myproject-123456", @@ -70,8 +70,8 @@ def test_from_dict(_mock_model_class, _mock_vertexai): } -@patch("google_vertex_haystack.generators.image_generator.vertexai") -@patch("google_vertex_haystack.generators.image_generator.ImageGenerationModel") +@patch("haystack_integrations.components.generators.google_vertex.image_generator.vertexai") +@patch("haystack_integrations.components.generators.google_vertex.image_generator.ImageGenerationModel") def test_run_calls_generate_images(mock_model_class, _mock_vertexai): mock_model = Mock() mock_model.generate_images.return_value = ImageGenerationResponse(images=[]) diff --git a/integrations/google_vertex/tests/test_question_answering.py b/integrations/google_vertex/tests/test_question_answering.py index 3495afcb2..3f414f0e0 100644 --- a/integrations/google_vertex/tests/test_question_answering.py +++ b/integrations/google_vertex/tests/test_question_answering.py @@ -2,11 +2,11 @@ from haystack.dataclasses.byte_stream import ByteStream -from google_vertex_haystack.generators.question_answering import VertexAIImageQA +from haystack_integrations.components.generators.google_vertex import VertexAIImageQA -@patch("google_vertex_haystack.generators.question_answering.vertexai") -@patch("google_vertex_haystack.generators.question_answering.ImageTextModel") +@patch("haystack_integrations.components.generators.google_vertex.question_answering.vertexai") +@patch("haystack_integrations.components.generators.google_vertex.question_answering.ImageTextModel") def test_init(mock_model_class, mock_vertexai): generator = VertexAIImageQA( model="imagetext", @@ -21,8 +21,8 @@ def test_init(mock_model_class, mock_vertexai): assert generator._kwargs == {"number_of_results": 3} -@patch("google_vertex_haystack.generators.question_answering.vertexai") -@patch("google_vertex_haystack.generators.question_answering.ImageTextModel") +@patch("haystack_integrations.components.generators.google_vertex.question_answering.vertexai") +@patch("haystack_integrations.components.generators.google_vertex.question_answering.ImageTextModel") def test_to_dict(_mock_model_class, _mock_vertexai): generator = VertexAIImageQA( model="imagetext", @@ -30,7 +30,7 @@ def test_to_dict(_mock_model_class, _mock_vertexai): number_of_results=3, ) assert generator.to_dict() == { - "type": "google_vertex_haystack.generators.question_answering.VertexAIImageQA", + "type": "haystack_integrations.components.generators.google_vertex.question_answering.VertexAIImageQA", "init_parameters": { "model": "imagetext", "project_id": "myproject-123456", @@ -40,12 +40,12 @@ def test_to_dict(_mock_model_class, _mock_vertexai): } -@patch("google_vertex_haystack.generators.question_answering.vertexai") -@patch("google_vertex_haystack.generators.question_answering.ImageTextModel") +@patch("haystack_integrations.components.generators.google_vertex.question_answering.vertexai") +@patch("haystack_integrations.components.generators.google_vertex.question_answering.ImageTextModel") def test_from_dict(_mock_model_class, _mock_vertexai): generator = VertexAIImageQA.from_dict( { - "type": "google_vertex_haystack.generators.question_answering.VertexAIImageQA", + "type": "haystack_integrations.components.generators.google_vertex.question_answering.VertexAIImageQA", "init_parameters": { "model": "imagetext", "project_id": "myproject-123456", @@ -60,8 +60,8 @@ def test_from_dict(_mock_model_class, _mock_vertexai): assert generator._kwargs == {"number_of_results": 3} -@patch("google_vertex_haystack.generators.question_answering.vertexai") -@patch("google_vertex_haystack.generators.question_answering.ImageTextModel") +@patch("haystack_integrations.components.generators.google_vertex.question_answering.vertexai") +@patch("haystack_integrations.components.generators.google_vertex.question_answering.ImageTextModel") def test_run_calls_ask_question(mock_model_class, _mock_vertexai): mock_model = Mock() mock_model.ask_question.return_value = [] diff --git a/integrations/google_vertex/tests/test_text_generator.py b/integrations/google_vertex/tests/test_text_generator.py index f2edbfc3b..3e5248dc7 100644 --- a/integrations/google_vertex/tests/test_text_generator.py +++ b/integrations/google_vertex/tests/test_text_generator.py @@ -2,11 +2,11 @@ from vertexai.language_models import GroundingSource -from google_vertex_haystack.generators.text_generator import VertexAITextGenerator +from haystack_integrations.components.generators.google_vertex import VertexAITextGenerator -@patch("google_vertex_haystack.generators.text_generator.vertexai") -@patch("google_vertex_haystack.generators.text_generator.TextGenerationModel") +@patch("haystack_integrations.components.generators.google_vertex.text_generator.vertexai") +@patch("haystack_integrations.components.generators.google_vertex.text_generator.TextGenerationModel") def test_init(mock_model_class, mock_vertexai): grounding_source = GroundingSource.VertexAISearch("1234", "us-central-1") generator = VertexAITextGenerator( @@ -20,15 +20,15 @@ def test_init(mock_model_class, mock_vertexai): assert generator._kwargs == {"temperature": 0.2, "grounding_source": grounding_source} -@patch("google_vertex_haystack.generators.text_generator.vertexai") -@patch("google_vertex_haystack.generators.text_generator.TextGenerationModel") +@patch("haystack_integrations.components.generators.google_vertex.text_generator.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 ) assert generator.to_dict() == { - "type": "google_vertex_haystack.generators.text_generator.VertexAITextGenerator", + "type": "haystack_integrations.components.generators.google_vertex.text_generator.VertexAITextGenerator", "init_parameters": { "model": "text-bison", "project_id": "myproject-123456", @@ -47,12 +47,12 @@ def test_to_dict(_mock_model_class, _mock_vertexai): } -@patch("google_vertex_haystack.generators.text_generator.vertexai") -@patch("google_vertex_haystack.generators.text_generator.TextGenerationModel") +@patch("haystack_integrations.components.generators.google_vertex.text_generator.vertexai") +@patch("haystack_integrations.components.generators.google_vertex.text_generator.TextGenerationModel") def test_from_dict(_mock_model_class, _mock_vertexai): generator = VertexAITextGenerator.from_dict( { - "type": "google_vertex_haystack.generators.text_generator.VertexAITextGenerator", + "type": "haystack_integrations.components.generators.google_vertex.text_generator.VertexAITextGenerator", "init_parameters": { "model": "text-bison", "project_id": "myproject-123456", @@ -79,8 +79,8 @@ def test_from_dict(_mock_model_class, _mock_vertexai): } -@patch("google_vertex_haystack.generators.text_generator.vertexai") -@patch("google_vertex_haystack.generators.text_generator.TextGenerationModel") +@patch("haystack_integrations.components.generators.google_vertex.text_generator.vertexai") +@patch("haystack_integrations.components.generators.google_vertex.text_generator.TextGenerationModel") def test_run_calls_get_captions(mock_model_class, _mock_vertexai): mock_model = Mock() mock_model.predict.return_value = MagicMock()