From 09f10f9cfb307c34e3365e706210111adc575459 Mon Sep 17 00:00:00 2001 From: "David S. Batista" Date: Tue, 13 Feb 2024 12:00:20 +0100 Subject: [PATCH 01/13] initial import --- integrations/amazon_bedrock/pyproject.toml | 4 +++- .../components/generators/amazon_bedrock/generator.py | 11 ++++++----- 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/integrations/amazon_bedrock/pyproject.toml b/integrations/amazon_bedrock/pyproject.toml index 8527d27a1..6c44c22a0 100644 --- a/integrations/amazon_bedrock/pyproject.toml +++ b/integrations/amazon_bedrock/pyproject.toml @@ -135,7 +135,9 @@ ignore = [ # Ignore complexity "C901", "PLR0911", "PLR0912", "PLR0913", "PLR0915", # Ignore unused params - "ARG001", "ARG002", "ARG005" + "ARG001", "ARG002", "ARG005", + # Ignore perform the call within the function + "B008" ] unfixable = [ # Don't touch unused imports diff --git a/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py b/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py index 8e89dab59..3eca4310a 100644 --- a/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py +++ b/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py @@ -6,6 +6,7 @@ import boto3 from botocore.exceptions import BotoCoreError, ClientError from haystack import component, default_from_dict, default_to_dict +from haystack.utils.auth import EnvVarSecret, Secret from .adapters import ( AI21LabsJurassic2Adapter, @@ -72,11 +73,11 @@ class AmazonBedrockGenerator: def __init__( self, model: str, - aws_access_key_id: Optional[str] = None, - aws_secret_access_key: Optional[str] = None, - aws_session_token: Optional[str] = None, - aws_region_name: Optional[str] = None, - aws_profile_name: Optional[str] = None, + aws_access_key_id: Optional[Secret] = EnvVarSecret.from_env_var("AWS_ACCESS_KEY_ID"), + aws_secret_access_key: Optional[Secret] = EnvVarSecret.from_env_var("AWS_SECRET_ACCESS_KEY"), + aws_session_token: Optional[str] = EnvVarSecret.from_env_var("AWS_SESSION_TOKEN"), + aws_region_name: Optional[Secret] = EnvVarSecret.from_env_var("AWS_DEFAULT_REGION"), + aws_profile_name: Optional[str] = EnvVarSecret.from_env_var("AWS_PROFILE"), max_length: Optional[int] = 100, **kwargs, ): From 41231b2a1844cca3f351540d0d254b61821822fa Mon Sep 17 00:00:00 2001 From: "David S. Batista" Date: Tue, 13 Feb 2024 14:31:20 +0100 Subject: [PATCH 02/13] initial import --- integrations/amazon_bedrock/pyproject.toml | 7 ++++++- .../components/generators/amazon_bedrock/generator.py | 11 +++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/integrations/amazon_bedrock/pyproject.toml b/integrations/amazon_bedrock/pyproject.toml index 6c44c22a0..105e19cd5 100644 --- a/integrations/amazon_bedrock/pyproject.toml +++ b/integrations/amazon_bedrock/pyproject.toml @@ -83,11 +83,13 @@ style = [ "ruff {args:.}", "black --check --diff {args:.}", ] + fmt = [ "black {args:.}", "ruff --fix {args:.}", "style", ] + all = [ "style", "typing", @@ -159,17 +161,20 @@ source_pkgs = ["src", "tests"] branch = true parallel = true - [tool.coverage.paths] amazon_bedrock_haystack = ["src/*"] tests = ["tests"] [tool.coverage.report] +omit = ["*__init__*"] exclude_lines = [ "no cov", "if __name__ == .__main__.:", "if TYPE_CHECKING:", ] +show_missing = true +precision = 2 + [[tool.mypy.overrides]] module = [ diff --git a/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py b/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py index 3eca4310a..4e94ed423 100644 --- a/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py +++ b/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py @@ -87,6 +87,17 @@ def __init__( self.model = model self.max_length = max_length + """ + if isinstance(aws_access_key_id, Secret): + aws_access_key_id = aws_access_key_id.resolve_value() + if isinstance(aws_secret_access_key, Secret): + aws_secret_access_key = aws_secret_access_key.resolve_value() + if isinstance(aws_session_token, Secret): + aws_session_token = aws_session_token.resolve_value() + if isinstance(aws_region_name, Secret): + aws_region_name = aws_region_name.resolve_value() + """ + try: session = self.get_aws_session( aws_access_key_id=aws_access_key_id, From cc29748f23d7cd2ef6086b8a58cb32ec5cf044c4 Mon Sep 17 00:00:00 2001 From: "David S. Batista" Date: Wed, 14 Feb 2024 10:13:19 +0100 Subject: [PATCH 03/13] wip --- .../generators/amazon_bedrock/generator.py | 16 ++-------------- 1 file changed, 2 insertions(+), 14 deletions(-) diff --git a/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py b/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py index 4e94ed423..f8ae48cb9 100644 --- a/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py +++ b/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py @@ -75,7 +75,7 @@ def __init__( model: str, aws_access_key_id: Optional[Secret] = EnvVarSecret.from_env_var("AWS_ACCESS_KEY_ID"), aws_secret_access_key: Optional[Secret] = EnvVarSecret.from_env_var("AWS_SECRET_ACCESS_KEY"), - aws_session_token: Optional[str] = EnvVarSecret.from_env_var("AWS_SESSION_TOKEN"), + aws_session_token: Optional[Secret] = EnvVarSecret.from_env_var("AWS_SESSION_TOKEN"), aws_region_name: Optional[Secret] = EnvVarSecret.from_env_var("AWS_DEFAULT_REGION"), aws_profile_name: Optional[str] = EnvVarSecret.from_env_var("AWS_PROFILE"), max_length: Optional[int] = 100, @@ -87,17 +87,6 @@ def __init__( self.model = model self.max_length = max_length - """ - if isinstance(aws_access_key_id, Secret): - aws_access_key_id = aws_access_key_id.resolve_value() - if isinstance(aws_secret_access_key, Secret): - aws_secret_access_key = aws_secret_access_key.resolve_value() - if isinstance(aws_session_token, Secret): - aws_session_token = aws_session_token.resolve_value() - if isinstance(aws_region_name, Secret): - aws_region_name = aws_region_name.resolve_value() - """ - try: session = self.get_aws_session( aws_access_key_id=aws_access_key_id, @@ -115,8 +104,7 @@ def __init__( raise AmazonBedrockConfigurationError(msg) from exception model_input_kwargs = kwargs - # We pop the model_max_length as it is not sent to the model - # but used to truncate the prompt if needed + # We pop the model_max_length as it is not sent to the model but used to truncate the prompt if needed model_max_length = kwargs.get("model_max_length", 4096) # Truncate prompt if prompt tokens > model_max_length-max_length From abd0bb4aebc5d171b0151d72e22d9c9f08eb6472 Mon Sep 17 00:00:00 2001 From: "David S. Batista" Date: Wed, 14 Feb 2024 14:13:07 +0100 Subject: [PATCH 04/13] wip: adding test coverage --- integrations/amazon_bedrock/pyproject.toml | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/integrations/amazon_bedrock/pyproject.toml b/integrations/amazon_bedrock/pyproject.toml index 105e19cd5..8ddd2f09a 100644 --- a/integrations/amazon_bedrock/pyproject.toml +++ b/integrations/amazon_bedrock/pyproject.toml @@ -157,24 +157,18 @@ ban-relative-imports = "parents" "tests/**/*" = ["PLR2004", "S101", "TID252"] [tool.coverage.run] -source_pkgs = ["src", "tests"] +source = ["haystack_integrations"] branch = true parallel = true -[tool.coverage.paths] -amazon_bedrock_haystack = ["src/*"] -tests = ["tests"] - [tool.coverage.report] -omit = ["*__init__*"] +omit = ["*/tests/*", "*/__init__.py"] +show_missing=true exclude_lines = [ "no cov", "if __name__ == .__main__.:", "if TYPE_CHECKING:", ] -show_missing = true -precision = 2 - [[tool.mypy.overrides]] module = [ From 33949aa294d434ad80900690ebe239456432a5dd Mon Sep 17 00:00:00 2001 From: "David S. Batista" Date: Wed, 14 Feb 2024 17:35:52 +0100 Subject: [PATCH 05/13] addin Secret and fixing tests --- .../generators/amazon_bedrock/generator.py | 39 ++++++++++++---- integrations/amazon_bedrock/tests/conftest.py | 34 ++++++++++++++ .../tests/test_amazon_bedrock.py | 46 ++++--------------- 3 files changed, 73 insertions(+), 46 deletions(-) create mode 100644 integrations/amazon_bedrock/tests/conftest.py diff --git a/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py b/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py index f8ae48cb9..7a75e58bb 100644 --- a/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py +++ b/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py @@ -73,11 +73,21 @@ class AmazonBedrockGenerator: def __init__( self, model: str, - aws_access_key_id: Optional[Secret] = EnvVarSecret.from_env_var("AWS_ACCESS_KEY_ID"), - aws_secret_access_key: Optional[Secret] = EnvVarSecret.from_env_var("AWS_SECRET_ACCESS_KEY"), - aws_session_token: Optional[Secret] = EnvVarSecret.from_env_var("AWS_SESSION_TOKEN"), - aws_region_name: Optional[Secret] = EnvVarSecret.from_env_var("AWS_DEFAULT_REGION"), - aws_profile_name: Optional[str] = EnvVarSecret.from_env_var("AWS_PROFILE"), + + aws_access_key_id: Optional[Secret] = EnvVarSecret( + env_vars=["AWS_ACCESS_KEY_ID"], strict=False).from_env_var("AWS_ACCESS_KEY_ID"), + + aws_secret_access_key: Optional[Secret] = EnvVarSecret( + env_vars=["AWS_SECRET_ACCESS_KEY"], strict=False).from_env_var("AWS_SECRET_ACCESS_KEY"), + + aws_session_token: Optional[Secret] = EnvVarSecret( + env_vars=["AWS_SESSION_TOKEN"], strict=False).from_env_var("AWS_SESSION_TOKEN"), + + aws_region_name: Optional[Secret] = EnvVarSecret( + env_vars=["AWS_DEFAULT_REGION"], strict=False).from_env_var("AWS_DEFAULT_REGION"), + + aws_profile_name: Optional[str] = EnvVarSecret( + env_vars=["AWS_PROFILE"], strict=False).from_env_var("AWS_PROFILE"), max_length: Optional[int] = 100, **kwargs, ): @@ -87,14 +97,23 @@ def __init__( self.model = model self.max_length = max_length + def resolve_secret(secret: Optional[Secret]) -> Optional[str]: + print(secret, secret.resolve_value()) + return secret.resolve_value() if secret else None + try: session = self.get_aws_session( - aws_access_key_id=aws_access_key_id, - aws_secret_access_key=aws_secret_access_key, - aws_session_token=aws_session_token, - aws_region_name=aws_region_name, - aws_profile_name=aws_profile_name, + aws_access_key_id=resolve_secret(aws_access_key_id), + aws_secret_access_key=resolve_secret(aws_secret_access_key), + aws_session_token=resolve_secret(aws_session_token), + aws_region_name=resolve_secret(aws_region_name), + aws_profile_name=resolve_secret(aws_profile_name), ) + + creds = session.get_credentials() + print("\n\n\n\nCreds") + print(creds.access_key, creds.secret_key, creds.token, creds.method, creds.expired, creds._refresh_needed) + self.client = session.client("bedrock-runtime") except Exception as exception: msg = ( diff --git a/integrations/amazon_bedrock/tests/conftest.py b/integrations/amazon_bedrock/tests/conftest.py new file mode 100644 index 000000000..3994aea6e --- /dev/null +++ b/integrations/amazon_bedrock/tests/conftest.py @@ -0,0 +1,34 @@ +import pytest +from unittest.mock import MagicMock, patch + + +@pytest.fixture +def set_env_variables(monkeypatch): + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "some_fake_id") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "some_fake_key") + monkeypatch.setenv("AWS_SESSION_TOKEN", "some_fake_token") + monkeypatch.setenv("AWS_DEFAULT_REGION", "fake_region") + monkeypatch.setenv("AWS_PROFILE", "some_fake_profile") + + +@pytest.fixture +def mock_auto_tokenizer(): + with patch("transformers.AutoTokenizer.from_pretrained", autospec=True) as mock_from_pretrained: + mock_tokenizer = MagicMock() + mock_from_pretrained.return_value = mock_tokenizer + yield mock_tokenizer + + +# create a fixture with mocked boto3 client and session +@pytest.fixture +def mock_boto3_session(): + with patch("boto3.Session") as mock_client: + yield mock_client + + +@pytest.fixture +def mock_prompt_handler(): + with patch( + "haystack_integrations.components.generators.amazon_bedrock.handlers.DefaultPromptHandler" + ) as mock_prompt_handler: + yield mock_prompt_handler diff --git a/integrations/amazon_bedrock/tests/test_amazon_bedrock.py b/integrations/amazon_bedrock/tests/test_amazon_bedrock.py index b08e9dfd5..2fdeb09b6 100644 --- a/integrations/amazon_bedrock/tests/test_amazon_bedrock.py +++ b/integrations/amazon_bedrock/tests/test_amazon_bedrock.py @@ -1,3 +1,4 @@ +import os from typing import Optional, Type from unittest.mock import MagicMock, call, patch @@ -16,42 +17,14 @@ from haystack_integrations.components.generators.amazon_bedrock.errors import AmazonBedrockConfigurationError -@pytest.fixture -def mock_auto_tokenizer(): - with patch("transformers.AutoTokenizer.from_pretrained", autospec=True) as mock_from_pretrained: - mock_tokenizer = MagicMock() - mock_from_pretrained.return_value = mock_tokenizer - yield mock_tokenizer - - -# create a fixture with mocked boto3 client and session -@pytest.fixture -def mock_boto3_session(): - with patch("boto3.Session") as mock_client: - yield mock_client - - -@pytest.fixture -def mock_prompt_handler(): - with patch( - "haystack_integrations.components.generators.amazon_bedrock.handlers.DefaultPromptHandler" - ) as mock_prompt_handler: - yield mock_prompt_handler - - @pytest.mark.unit -def test_to_dict(mock_auto_tokenizer, mock_boto3_session): +def test_to_dict(mock_auto_tokenizer, mock_boto3_session, set_env_variables): """ Test that the to_dict method returns the correct dictionary without aws credentials """ generator = AmazonBedrockGenerator( model="anthropic.claude-v2", max_length=99, - aws_access_key_id="some_fake_id", - aws_secret_access_key="some_fake_key", - aws_session_token="some_fake_token", - aws_profile_name="some_fake_profile", - aws_region_name="fake_region", ) expected_dict = { @@ -66,7 +39,7 @@ def test_to_dict(mock_auto_tokenizer, mock_boto3_session): @pytest.mark.unit -def test_from_dict(mock_auto_tokenizer, mock_boto3_session): +def test_from_dict(mock_auto_tokenizer, mock_boto3_session, set_env_variables): """ Test that the from_dict method returns the correct object """ @@ -90,14 +63,15 @@ def test_default_constructor(mock_auto_tokenizer, mock_boto3_session): Test that the default constructor sets the correct values """ + os.environ["AWS_ACCESS_KEY_ID"] = "some_fake_id" + os.environ["AWS_SECRET_ACCESS_KEY"] = "some_fake_key" + os.environ["AWS_SESSION_TOKEN"] = "some_fake_token" + os.environ["AWS_DEFAULT_REGION"] = "fake_region" + os.environ["AWS_PROFILE"] = "some_fake_profile" + layer = AmazonBedrockGenerator( model="anthropic.claude-v2", max_length=99, - aws_access_key_id="some_fake_id", - aws_secret_access_key="some_fake_key", - aws_session_token="some_fake_token", - aws_profile_name="some_fake_profile", - aws_region_name="fake_region", ) assert layer.max_length == 99 @@ -120,7 +94,7 @@ def test_default_constructor(mock_auto_tokenizer, mock_boto3_session): @pytest.mark.unit -def test_constructor_prompt_handler_initialized(mock_auto_tokenizer, mock_boto3_session): +def test_constructor_prompt_handler_initialized(mock_auto_tokenizer, mock_boto3_session, mock_prompt_handler): """ Test that the constructor sets the prompt_handler correctly, with the correct model_max_length for llama-2 """ From d6668d58189bde365365822eb55c9ce55b7472e4 Mon Sep 17 00:00:00 2001 From: "David S. Batista" Date: Wed, 14 Feb 2024 17:55:26 +0100 Subject: [PATCH 06/13] cleaning --- .../generators/amazon_bedrock/generator.py | 34 ++++++++----------- integrations/amazon_bedrock/tests/conftest.py | 3 +- 2 files changed, 16 insertions(+), 21 deletions(-) diff --git a/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py b/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py index 7a75e58bb..2517ccfb2 100644 --- a/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py +++ b/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py @@ -73,21 +73,21 @@ class AmazonBedrockGenerator: def __init__( self, model: str, - - aws_access_key_id: Optional[Secret] = EnvVarSecret( - env_vars=["AWS_ACCESS_KEY_ID"], strict=False).from_env_var("AWS_ACCESS_KEY_ID"), - + aws_access_key_id: Optional[Secret] = EnvVarSecret(env_vars=["AWS_ACCESS_KEY_ID"], strict=False).from_env_var( + "AWS_ACCESS_KEY_ID" + ), aws_secret_access_key: Optional[Secret] = EnvVarSecret( - env_vars=["AWS_SECRET_ACCESS_KEY"], strict=False).from_env_var("AWS_SECRET_ACCESS_KEY"), - - aws_session_token: Optional[Secret] = EnvVarSecret( - env_vars=["AWS_SESSION_TOKEN"], strict=False).from_env_var("AWS_SESSION_TOKEN"), - - aws_region_name: Optional[Secret] = EnvVarSecret( - env_vars=["AWS_DEFAULT_REGION"], strict=False).from_env_var("AWS_DEFAULT_REGION"), - - aws_profile_name: Optional[str] = EnvVarSecret( - env_vars=["AWS_PROFILE"], strict=False).from_env_var("AWS_PROFILE"), + env_vars=["AWS_SECRET_ACCESS_KEY"], strict=False + ).from_env_var("AWS_SECRET_ACCESS_KEY"), + aws_session_token: Optional[Secret] = EnvVarSecret(env_vars=["AWS_SESSION_TOKEN"], strict=False).from_env_var( + "AWS_SESSION_TOKEN" + ), + aws_region_name: Optional[Secret] = EnvVarSecret(env_vars=["AWS_DEFAULT_REGION"], strict=False).from_env_var( + "AWS_DEFAULT_REGION" + ), + aws_profile_name: Optional[str] = EnvVarSecret(env_vars=["AWS_PROFILE"], strict=False).from_env_var( + "AWS_PROFILE" + ), max_length: Optional[int] = 100, **kwargs, ): @@ -98,7 +98,6 @@ def __init__( self.max_length = max_length def resolve_secret(secret: Optional[Secret]) -> Optional[str]: - print(secret, secret.resolve_value()) return secret.resolve_value() if secret else None try: @@ -109,11 +108,6 @@ def resolve_secret(secret: Optional[Secret]) -> Optional[str]: aws_region_name=resolve_secret(aws_region_name), aws_profile_name=resolve_secret(aws_profile_name), ) - - creds = session.get_credentials() - print("\n\n\n\nCreds") - print(creds.access_key, creds.secret_key, creds.token, creds.method, creds.expired, creds._refresh_needed) - self.client = session.client("bedrock-runtime") except Exception as exception: msg = ( diff --git a/integrations/amazon_bedrock/tests/conftest.py b/integrations/amazon_bedrock/tests/conftest.py index 3994aea6e..4c8ce688c 100644 --- a/integrations/amazon_bedrock/tests/conftest.py +++ b/integrations/amazon_bedrock/tests/conftest.py @@ -1,6 +1,7 @@ -import pytest from unittest.mock import MagicMock, patch +import pytest + @pytest.fixture def set_env_variables(monkeypatch): From 1b96111ac83d075b1815f1e30fc9406f941b5082 Mon Sep 17 00:00:00 2001 From: "David S. Batista" Date: Thu, 15 Feb 2024 09:47:26 +0100 Subject: [PATCH 07/13] using staticmethod directly --- .../generators/amazon_bedrock/generator.py | 22 +++++-------------- 1 file changed, 6 insertions(+), 16 deletions(-) diff --git a/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py b/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py index 2517ccfb2..102e5d535 100644 --- a/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py +++ b/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py @@ -6,7 +6,7 @@ import boto3 from botocore.exceptions import BotoCoreError, ClientError from haystack import component, default_from_dict, default_to_dict -from haystack.utils.auth import EnvVarSecret, Secret +from haystack.utils.auth import Secret from .adapters import ( AI21LabsJurassic2Adapter, @@ -73,21 +73,11 @@ class AmazonBedrockGenerator: def __init__( self, model: str, - aws_access_key_id: Optional[Secret] = EnvVarSecret(env_vars=["AWS_ACCESS_KEY_ID"], strict=False).from_env_var( - "AWS_ACCESS_KEY_ID" - ), - aws_secret_access_key: Optional[Secret] = EnvVarSecret( - env_vars=["AWS_SECRET_ACCESS_KEY"], strict=False - ).from_env_var("AWS_SECRET_ACCESS_KEY"), - aws_session_token: Optional[Secret] = EnvVarSecret(env_vars=["AWS_SESSION_TOKEN"], strict=False).from_env_var( - "AWS_SESSION_TOKEN" - ), - aws_region_name: Optional[Secret] = EnvVarSecret(env_vars=["AWS_DEFAULT_REGION"], strict=False).from_env_var( - "AWS_DEFAULT_REGION" - ), - aws_profile_name: Optional[str] = EnvVarSecret(env_vars=["AWS_PROFILE"], strict=False).from_env_var( - "AWS_PROFILE" - ), + aws_access_key_id: Optional[Secret] = Secret.from_env_var(["AWS_ACCESS_KEY_ID"], strict=False), + aws_secret_access_key: Optional[Secret] = Secret.from_env_var(["AWS_SECRET_ACCESS_KEY"], strict=False), + aws_session_token: Optional[Secret] = Secret.from_env_var(["AWS_SESSION_TOKEN"], strict=False), + aws_region_name: Optional[Secret] = Secret.from_env_var(["AWS_DEFAULT_REGION"], strict=False), + aws_profile_name: Optional[Secret] = Secret.from_env_var(["AWS_PROFILE"], strict=False), max_length: Optional[int] = 100, **kwargs, ): From 87dcfdb69ceba1236cf98194ffd8c7a1b785d518 Mon Sep 17 00:00:00 2001 From: "David S. Batista" Date: Thu, 15 Feb 2024 10:08:40 +0100 Subject: [PATCH 08/13] removing ignore from B008 config and setting up in-line --- integrations/amazon_bedrock/pyproject.toml | 2 -- .../generators/amazon_bedrock/generator.py | 12 +++++++----- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/integrations/amazon_bedrock/pyproject.toml b/integrations/amazon_bedrock/pyproject.toml index 8ddd2f09a..8366c9449 100644 --- a/integrations/amazon_bedrock/pyproject.toml +++ b/integrations/amazon_bedrock/pyproject.toml @@ -138,8 +138,6 @@ ignore = [ "C901", "PLR0911", "PLR0912", "PLR0913", "PLR0915", # Ignore unused params "ARG001", "ARG002", "ARG005", - # Ignore perform the call within the function - "B008" ] unfixable = [ # Don't touch unused imports diff --git a/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py b/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py index 102e5d535..a9e6ad7cc 100644 --- a/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py +++ b/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py @@ -73,11 +73,13 @@ class AmazonBedrockGenerator: def __init__( self, model: str, - aws_access_key_id: Optional[Secret] = Secret.from_env_var(["AWS_ACCESS_KEY_ID"], strict=False), - aws_secret_access_key: Optional[Secret] = Secret.from_env_var(["AWS_SECRET_ACCESS_KEY"], strict=False), - aws_session_token: Optional[Secret] = Secret.from_env_var(["AWS_SESSION_TOKEN"], strict=False), - aws_region_name: Optional[Secret] = Secret.from_env_var(["AWS_DEFAULT_REGION"], strict=False), - aws_profile_name: Optional[Secret] = Secret.from_env_var(["AWS_PROFILE"], strict=False), + aws_access_key_id: Optional[Secret] = Secret.from_env_var(["AWS_ACCESS_KEY_ID"], strict=False), # noqa: B008 + aws_secret_access_key: Optional[Secret] = Secret.from_env_var( # noqa: B008 + ["AWS_SECRET_ACCESS_KEY"], strict=False + ), + aws_session_token: Optional[Secret] = Secret.from_env_var(["AWS_SESSION_TOKEN"], strict=False), # noqa: B008 + aws_region_name: Optional[Secret] = Secret.from_env_var(["AWS_DEFAULT_REGION"], strict=False), # noqa: B008 + aws_profile_name: Optional[Secret] = Secret.from_env_var(["AWS_PROFILE"], strict=False), # noqa: B008 max_length: Optional[int] = 100, **kwargs, ): From 97c83a5eedceb9d7935d4e4b40ceea05e1c1722a Mon Sep 17 00:00:00 2001 From: "David S. Batista" Date: Thu, 15 Feb 2024 10:15:09 +0100 Subject: [PATCH 09/13] addin Secret and fixing tests for chat component --- .../amazon_bedrock/chat/chat_generator.py | 26 ++++++++++++------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/chat/chat_generator.py b/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/chat/chat_generator.py index 804d44413..696f4a835 100644 --- a/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/chat/chat_generator.py +++ b/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/chat/chat_generator.py @@ -8,6 +8,7 @@ from haystack import component, default_from_dict, default_to_dict from haystack.components.generators.utils import deserialize_callback_handler from haystack.dataclasses import ChatMessage, StreamingChunk +from haystack.utils.auth import Secret from haystack_integrations.components.generators.amazon_bedrock.errors import ( AmazonBedrockConfigurationError, @@ -61,11 +62,13 @@ class AmazonBedrockChatGenerator: def __init__( self, model: str, - aws_access_key_id: Optional[str] = None, - aws_secret_access_key: Optional[str] = None, - aws_session_token: Optional[str] = None, - aws_region_name: Optional[str] = None, - aws_profile_name: Optional[str] = None, + aws_access_key_id: Optional[Secret] = Secret.from_env_var(["AWS_ACCESS_KEY_ID"], strict=False), # noqa: B008 + aws_secret_access_key: Optional[Secret] = Secret.from_env_var( # noqa: B008 + ["AWS_SECRET_ACCESS_KEY"], strict=False + ), + aws_session_token: Optional[Secret] = Secret.from_env_var(["AWS_SESSION_TOKEN"], strict=False), # noqa: B008 + aws_region_name: Optional[Secret] = Secret.from_env_var(["AWS_DEFAULT_REGION"], strict=False), # noqa: B008 + aws_profile_name: Optional[Secret] = Secret.from_env_var(["AWS_PROFILE"], strict=False), # noqa: B008 generation_kwargs: Optional[Dict[str, Any]] = None, stop_words: Optional[List[str]] = None, streaming_callback: Optional[Callable[[StreamingChunk], None]] = None, @@ -111,13 +114,16 @@ def __init__( self.model_adapter = model_adapter_cls(generation_kwargs or {}) # create the AWS session and client + def resolve_secret(secret: Optional[Secret]) -> Optional[str]: + return secret.resolve_value() if secret else None + try: session = self.get_aws_session( - aws_access_key_id=aws_access_key_id, - aws_secret_access_key=aws_secret_access_key, - aws_session_token=aws_session_token, - aws_region_name=aws_region_name, - aws_profile_name=aws_profile_name, + aws_access_key_id=resolve_secret(aws_access_key_id), + aws_secret_access_key=resolve_secret(aws_secret_access_key), + aws_session_token=resolve_secret(aws_session_token), + aws_region_name=resolve_secret(aws_region_name), + aws_profile_name=resolve_secret(aws_profile_name), ) self.client = session.client("bedrock-runtime") except Exception as exception: From 293eb2489f018be3a468a8542f4021531314190a Mon Sep 17 00:00:00 2001 From: "David S. Batista" Date: Thu, 15 Feb 2024 10:22:02 +0100 Subject: [PATCH 10/13] addin Secret and fixing tests for the chat component --- .../tests/test_amazon_chat_bedrock.py | 14 ++------------ 1 file changed, 2 insertions(+), 12 deletions(-) diff --git a/integrations/amazon_bedrock/tests/test_amazon_chat_bedrock.py b/integrations/amazon_bedrock/tests/test_amazon_chat_bedrock.py index 9592b5b39..6541c1b4f 100644 --- a/integrations/amazon_bedrock/tests/test_amazon_chat_bedrock.py +++ b/integrations/amazon_bedrock/tests/test_amazon_chat_bedrock.py @@ -38,17 +38,12 @@ def mock_prompt_handler(): yield mock_prompt_handler -def test_to_dict(mock_auto_tokenizer, mock_boto3_session): +def test_to_dict(mock_auto_tokenizer, mock_boto3_session, set_env_variables): """ Test that the to_dict method returns the correct dictionary without aws credentials """ generator = AmazonBedrockChatGenerator( model="anthropic.claude-v2", - aws_access_key_id="some_fake_id", - aws_secret_access_key="some_fake_key", - aws_session_token="some_fake_token", - aws_profile_name="some_fake_profile", - aws_region_name="fake_region", generation_kwargs={"temperature": 0.7}, streaming_callback=print_streaming_chunk, ) @@ -84,18 +79,13 @@ def test_from_dict(mock_auto_tokenizer, mock_boto3_session): assert generator.streaming_callback == print_streaming_chunk -def test_default_constructor(mock_auto_tokenizer, mock_boto3_session): +def test_default_constructor(mock_auto_tokenizer, mock_boto3_session, set_env_variables): """ Test that the default constructor sets the correct values """ layer = AmazonBedrockChatGenerator( model="anthropic.claude-v2", - aws_access_key_id="some_fake_id", - aws_secret_access_key="some_fake_key", - aws_session_token="some_fake_token", - aws_profile_name="some_fake_profile", - aws_region_name="fake_region", ) assert layer.model == "anthropic.claude-v2" From 3f2807bd196fc2c2d5e75017faaaf8a2e7f81b51 Mon Sep 17 00:00:00 2001 From: "David S. Batista" Date: Thu, 15 Feb 2024 11:59:39 +0100 Subject: [PATCH 11/13] wip --- .../components/generators/amazon_bedrock/generator.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py b/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py index a9e6ad7cc..253048b9c 100644 --- a/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py +++ b/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py @@ -73,13 +73,13 @@ class AmazonBedrockGenerator: def __init__( self, model: str, - aws_access_key_id: Optional[Secret] = Secret.from_env_var(["AWS_ACCESS_KEY_ID"], strict=False), # noqa: B008 + aws_access_key_id: Optional[Secret] = Secret.from_env_var("AWS_ACCESS_KEY_ID", strict=False), # noqa: B008 aws_secret_access_key: Optional[Secret] = Secret.from_env_var( # noqa: B008 - ["AWS_SECRET_ACCESS_KEY"], strict=False + "AWS_SECRET_ACCESS_KEY", strict=False ), - aws_session_token: Optional[Secret] = Secret.from_env_var(["AWS_SESSION_TOKEN"], strict=False), # noqa: B008 - aws_region_name: Optional[Secret] = Secret.from_env_var(["AWS_DEFAULT_REGION"], strict=False), # noqa: B008 - aws_profile_name: Optional[Secret] = Secret.from_env_var(["AWS_PROFILE"], strict=False), # noqa: B008 + aws_session_token: Optional[Secret] = Secret.from_env_var("AWS_SESSION_TOKEN", strict=False), # noqa: B008 + aws_region_name: Optional[Secret] = Secret.from_env_var("AWS_DEFAULT_REGION", strict=False), # noqa: B008 + aws_profile_name: Optional[Secret] = Secret.from_env_var("AWS_PROFILE", strict=False), # noqa: B008 max_length: Optional[int] = 100, **kwargs, ): From 75b3d64338398cbfb5a7147c0daf8026db8bd2cd Mon Sep 17 00:00:00 2001 From: "David S. Batista" Date: Thu, 15 Feb 2024 18:10:41 +0100 Subject: [PATCH 12/13] fixing to_dict from_dict --- .../generators/amazon_bedrock/generator.py | 16 +++++++++++++++- .../tests/test_amazon_bedrock.py | 19 +++++++++++-------- 2 files changed, 26 insertions(+), 9 deletions(-) diff --git a/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py b/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py index 253048b9c..1a8bb04f0 100644 --- a/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py +++ b/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/generator.py @@ -6,7 +6,7 @@ import boto3 from botocore.exceptions import BotoCoreError, ClientError from haystack import component, default_from_dict, default_to_dict -from haystack.utils.auth import Secret +from haystack.utils.auth import Secret, deserialize_secrets_inplace from .adapters import ( AI21LabsJurassic2Adapter, @@ -88,6 +88,11 @@ def __init__( raise ValueError(msg) self.model = model self.max_length = max_length + self.aws_access_key_id = aws_access_key_id + self.aws_secret_access_key = aws_secret_access_key + self.aws_session_token = aws_session_token + self.aws_region_name = aws_region_name + self.aws_profile_name = aws_profile_name def resolve_secret(secret: Optional[Secret]) -> Optional[str]: return secret.resolve_value() if secret else None @@ -303,6 +308,11 @@ def to_dict(self) -> Dict[str, Any]: """ return default_to_dict( self, + aws_access_key_id=self.aws_access_key_id.to_dict() if self.aws_access_key_id else None, + aws_secret_access_key=self.aws_secret_access_key.to_dict() if self.aws_secret_access_key else None, + aws_session_token=self.aws_session_token.to_dict() if self.aws_session_token else None, + aws_region_name=self.aws_region_name.to_dict() if self.aws_region_name else None, + aws_profile_name=self.aws_profile_name.to_dict() if self.aws_profile_name else None, model=self.model, max_length=self.max_length, ) @@ -314,4 +324,8 @@ def from_dict(cls, data: Dict[str, Any]) -> "AmazonBedrockGenerator": :param data: The dictionary representation of this component. :return: The deserialized component instance. """ + deserialize_secrets_inplace( + data["init_parameters"], + ["aws_access_key_id", "aws_secret_access_key", "aws_session_token", "aws_region_name", "aws_profile_name"], + ) return default_from_dict(cls, data) diff --git a/integrations/amazon_bedrock/tests/test_amazon_bedrock.py b/integrations/amazon_bedrock/tests/test_amazon_bedrock.py index 2fdeb09b6..f4d10a9b2 100644 --- a/integrations/amazon_bedrock/tests/test_amazon_bedrock.py +++ b/integrations/amazon_bedrock/tests/test_amazon_bedrock.py @@ -1,4 +1,3 @@ -import os from typing import Optional, Type from unittest.mock import MagicMock, call, patch @@ -30,6 +29,11 @@ def test_to_dict(mock_auto_tokenizer, mock_boto3_session, set_env_variables): expected_dict = { "type": "haystack_integrations.components.generators.amazon_bedrock.generator.AmazonBedrockGenerator", "init_parameters": { + "aws_access_key_id": {"type": "env_var", "env_vars": ["AWS_ACCESS_KEY_ID"], "strict": False}, + "aws_secret_access_key": {"type": "env_var", "env_vars": ["AWS_SECRET_ACCESS_KEY"], "strict": False}, + "aws_session_token": {"type": "env_var", "env_vars": ["AWS_SESSION_TOKEN"], "strict": False}, + "aws_region_name": {"type": "env_var", "env_vars": ["AWS_DEFAULT_REGION"], "strict": False}, + "aws_profile_name": {"type": "env_var", "env_vars": ["AWS_PROFILE"], "strict": False}, "model": "anthropic.claude-v2", "max_length": 99, }, @@ -47,6 +51,11 @@ def test_from_dict(mock_auto_tokenizer, mock_boto3_session, set_env_variables): { "type": "haystack_integrations.components.generators.amazon_bedrock.generator.AmazonBedrockGenerator", "init_parameters": { + "aws_access_key_id": {"type": "env_var", "env_vars": ["AWS_ACCESS_KEY_ID"], "strict": False}, + "aws_secret_access_key": {"type": "env_var", "env_vars": ["AWS_SECRET_ACCESS_KEY"], "strict": False}, + "aws_session_token": {"type": "env_var", "env_vars": ["AWS_SESSION_TOKEN"], "strict": False}, + "aws_region_name": {"type": "env_var", "env_vars": ["AWS_DEFAULT_REGION"], "strict": False}, + "aws_profile_name": {"type": "env_var", "env_vars": ["AWS_PROFILE"], "strict": False}, "model": "anthropic.claude-v2", "max_length": 99, }, @@ -58,17 +67,11 @@ def test_from_dict(mock_auto_tokenizer, mock_boto3_session, set_env_variables): @pytest.mark.unit -def test_default_constructor(mock_auto_tokenizer, mock_boto3_session): +def test_default_constructor(mock_auto_tokenizer, mock_boto3_session, set_env_variables): """ Test that the default constructor sets the correct values """ - os.environ["AWS_ACCESS_KEY_ID"] = "some_fake_id" - os.environ["AWS_SECRET_ACCESS_KEY"] = "some_fake_key" - os.environ["AWS_SESSION_TOKEN"] = "some_fake_token" - os.environ["AWS_DEFAULT_REGION"] = "fake_region" - os.environ["AWS_PROFILE"] = "some_fake_profile" - layer = AmazonBedrockGenerator( model="anthropic.claude-v2", max_length=99, From ac0b3c2eb8ff41237e3543057c88446f84f552a1 Mon Sep 17 00:00:00 2001 From: "David S. Batista" Date: Thu, 15 Feb 2024 18:18:10 +0100 Subject: [PATCH 13/13] fixing to_dict from_dict to the chat component as well --- .../amazon_bedrock/chat/chat_generator.py | 16 +++++++++++++++- .../tests/test_amazon_chat_bedrock.py | 10 ++++++++++ 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/chat/chat_generator.py b/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/chat/chat_generator.py index 696f4a835..6ce671e68 100644 --- a/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/chat/chat_generator.py +++ b/integrations/amazon_bedrock/src/haystack_integrations/components/generators/amazon_bedrock/chat/chat_generator.py @@ -8,7 +8,7 @@ from haystack import component, default_from_dict, default_to_dict from haystack.components.generators.utils import deserialize_callback_handler from haystack.dataclasses import ChatMessage, StreamingChunk -from haystack.utils.auth import Secret +from haystack.utils.auth import Secret, deserialize_secrets_inplace from haystack_integrations.components.generators.amazon_bedrock.errors import ( AmazonBedrockConfigurationError, @@ -105,6 +105,11 @@ def __init__( msg = "'model' cannot be None or empty string" raise ValueError(msg) self.model = model + self.aws_access_key_id = aws_access_key_id + self.aws_secret_access_key = aws_secret_access_key + self.aws_session_token = aws_session_token + self.aws_region_name = aws_region_name + self.aws_profile_name = aws_profile_name # get the model adapter for the given model model_adapter_cls = self.get_model_adapter(model=model) @@ -235,6 +240,11 @@ def to_dict(self) -> Dict[str, Any]: """ return default_to_dict( self, + aws_access_key_id=self.aws_access_key_id.to_dict() if self.aws_access_key_id else None, + aws_secret_access_key=self.aws_secret_access_key.to_dict() if self.aws_secret_access_key else None, + aws_session_token=self.aws_session_token.to_dict() if self.aws_session_token else None, + aws_region_name=self.aws_region_name.to_dict() if self.aws_region_name else None, + aws_profile_name=self.aws_profile_name.to_dict() if self.aws_profile_name else None, model=self.model, stop_words=self.stop_words, generation_kwargs=self.model_adapter.generation_kwargs, @@ -252,4 +262,8 @@ def from_dict(cls, data: Dict[str, Any]) -> "AmazonBedrockChatGenerator": serialized_callback_handler = init_params.get("streaming_callback") if serialized_callback_handler: data["init_parameters"]["streaming_callback"] = deserialize_callback_handler(serialized_callback_handler) + deserialize_secrets_inplace( + data["init_parameters"], + ["aws_access_key_id", "aws_secret_access_key", "aws_session_token", "aws_region_name", "aws_profile_name"], + ) return default_from_dict(cls, data) diff --git a/integrations/amazon_bedrock/tests/test_amazon_chat_bedrock.py b/integrations/amazon_bedrock/tests/test_amazon_chat_bedrock.py index 6541c1b4f..574aab5cc 100644 --- a/integrations/amazon_bedrock/tests/test_amazon_chat_bedrock.py +++ b/integrations/amazon_bedrock/tests/test_amazon_chat_bedrock.py @@ -50,6 +50,11 @@ def test_to_dict(mock_auto_tokenizer, mock_boto3_session, set_env_variables): expected_dict = { "type": clazz, "init_parameters": { + "aws_access_key_id": {"type": "env_var", "env_vars": ["AWS_ACCESS_KEY_ID"], "strict": False}, + "aws_secret_access_key": {"type": "env_var", "env_vars": ["AWS_SECRET_ACCESS_KEY"], "strict": False}, + "aws_session_token": {"type": "env_var", "env_vars": ["AWS_SESSION_TOKEN"], "strict": False}, + "aws_region_name": {"type": "env_var", "env_vars": ["AWS_DEFAULT_REGION"], "strict": False}, + "aws_profile_name": {"type": "env_var", "env_vars": ["AWS_PROFILE"], "strict": False}, "model": "anthropic.claude-v2", "generation_kwargs": {"temperature": 0.7}, "stop_words": [], @@ -68,6 +73,11 @@ def test_from_dict(mock_auto_tokenizer, mock_boto3_session): { "type": clazz, "init_parameters": { + "aws_access_key_id": {"type": "env_var", "env_vars": ["AWS_ACCESS_KEY_ID"], "strict": False}, + "aws_secret_access_key": {"type": "env_var", "env_vars": ["AWS_SECRET_ACCESS_KEY"], "strict": False}, + "aws_session_token": {"type": "env_var", "env_vars": ["AWS_SESSION_TOKEN"], "strict": False}, + "aws_region_name": {"type": "env_var", "env_vars": ["AWS_DEFAULT_REGION"], "strict": False}, + "aws_profile_name": {"type": "env_var", "env_vars": ["AWS_PROFILE"], "strict": False}, "model": "anthropic.claude-v2", "generation_kwargs": {"temperature": 0.7}, "streaming_callback": "haystack.components.generators.utils.print_streaming_chunk",