From 1d99bb5db668ea5c402017b405dc5a1271496ab0 Mon Sep 17 00:00:00 2001 From: Tao Chen Date: Mon, 3 Mar 2025 14:18:02 -0800 Subject: [PATCH 1/6] Improve Bedrock connector setup docs --- python/.cspell.json | 1 + python/samples/concepts/setup/ALL_SETTINGS.md | 8 ++++--- .../semantic_kernel/connectors/ai/README.md | 2 +- .../connectors/ai/bedrock/README.md | 21 ++++++++++++++++++ .../ai/bedrock/services/bedrock_base.py | 22 +++++++++++++++++++ .../services/bedrock_chat_completion.py | 6 ++--- .../services/bedrock_text_completion.py | 5 ++--- .../services/bedrock_text_embedding.py | 5 ++--- 8 files changed, 56 insertions(+), 14 deletions(-) diff --git a/python/.cspell.json b/python/.cspell.json index d236099b5943..1a99593354f1 100644 --- a/python/.cspell.json +++ b/python/.cspell.json @@ -51,6 +51,7 @@ "nopep", "NOSQL", "ollama", + "Onnx", "onyourdatatest", "OPENAI", "opentelemetry", diff --git a/python/samples/concepts/setup/ALL_SETTINGS.md b/python/samples/concepts/setup/ALL_SETTINGS.md index 1f2536ad4738..4deba2610935 100644 --- a/python/samples/concepts/setup/ALL_SETTINGS.md +++ b/python/samples/concepts/setup/ALL_SETTINGS.md @@ -1,4 +1,6 @@ -## AI Service Settings used across SK: +# Semantic Kernel Settings + +## AI Service Settings used across SK | Provider | Service | Constructor Settings | Environment Variable | Required? | Settings Class | | --- | --- | --- | --- | --- | --- | @@ -36,7 +38,7 @@ | Onnx | [OnnxGenAIChatCompletion](../../../semantic_kernel/connectors/ai/onnx/services/onnx_gen_ai_chat_completion.py) | template,
ai_model_path | N/A,
ONNX_GEN_AI_CHAT_MODEL_FOLDER | Yes,
Yes | [OnnxGenAISettings](../../../semantic_kernel/connectors/ai/onnx/onnx_gen_ai_settings.py) | | | [OnnxGenAITextCompletion](../../../semantic_kernel/connectors/ai/onnx/services/onnx_gen_ai_text_completion.py) | ai_model_path | ONNX_GEN_AI_TEXT_MODEL_FOLDER | Yes | | -## Memory Service Settings used across SK: +## Memory Service Settings used across SK | Provider | Service | Constructor Settings | Environment Variable | Required? | Settings Class | | --- | --- | --- | --- | --- | --- | @@ -49,7 +51,7 @@ | Redis | [RedisMemoryService](../../../semantic_kernel/connectors/memory/redis/redis_memory_store.py) | connection_string | REDIS_CONNECTION_STRING | Yes | [RedisSettings](../../../semantic_kernel/connectors/memory/redis/redis_settings.py) | | Weaviate | [WeaviateMemoryService](../../../semantic_kernel/connectors/memory/weaviate/weaviate_memory_store.py) | url,
api_key,
use_embed | WEAVIATE_URL,
WEAVIATE_API_KEY,
WEAVIATE_USE_EMBED | No,
No,
No | [WeaviateSettings](../../../semantic_kernel/connectors/memory/weaviate/weaviate_settings.py) | -## Other settings used: +## Other settings used | Provider | Service | Constructor Settings | Environment Variable | Required? | Settings Class | | --- | --- | --- | --- | --- | --- | diff --git a/python/semantic_kernel/connectors/ai/README.md b/python/semantic_kernel/connectors/ai/README.md index 997a33427c65..bad47044b835 100644 --- a/python/semantic_kernel/connectors/ai/README.md +++ b/python/semantic_kernel/connectors/ai/README.md @@ -3,6 +3,7 @@ This directory contains the implementation of the AI connectors (aka AI services) that are used to interact with AI models. Depending on the modality, the AI connector can inherit from one of the following classes: + - [`ChatCompletionClientBase`](./chat_completion_client_base.py) for chat completion tasks. - [`TextCompletionClientBase`](./text_completion_client_base.py) for text completion tasks. - [`AudioToTextClientBase`](./audio_to_text_client_base.py) for audio to text tasks. @@ -10,7 +11,6 @@ Depending on the modality, the AI connector can inherit from one of the followin - [`TextToImageClientBase`](./text_to_image_client_base.py) for text to image tasks. - [`EmbeddingGeneratorBase`](./embeddings/embedding_generator_base.py) for text embedding tasks. - All base clients inherit from the [`AIServiceClientBase`](../../services/ai_service_client_base.py) class. ## Existing AI connectors diff --git a/python/semantic_kernel/connectors/ai/bedrock/README.md b/python/semantic_kernel/connectors/ai/bedrock/README.md index 9f48879b54cb..6dee8f9c5bbb 100644 --- a/python/semantic_kernel/connectors/ai/bedrock/README.md +++ b/python/semantic_kernel/connectors/ai/bedrock/README.md @@ -11,6 +11,27 @@ Follow this [guide](https://boto3.amazonaws.com/v1/documentation/api/latest/guide/quickstart.html#configuration) to configure your environment to use the Bedrock API. +> Please configure the `aws_access_key_id`, `aws_secret_access_key`, and `region` otherwise you will need to create custom clients for the services. For example: + + ```python + runtime_client=boto.client( + "bedrock-runtime", + aws_access_key_id="your_access_key", + aws_secret_access_key="your_secret_key", + region_name="your_region", + [...other parameters you may need...] + ) + client=boto.client( + "bedrock", + aws_access_key_id="your_access_key", + aws_secret_access_key="your_secret_key", + region_name="your_region", + [...other parameters you may need...] + ) + + bedrock_chat_completion_service = BedrockChatCompletion(runtime_client=runtime_client, client=client) + ``` + ## Supports ### Region diff --git a/python/semantic_kernel/connectors/ai/bedrock/services/bedrock_base.py b/python/semantic_kernel/connectors/ai/bedrock/services/bedrock_base.py index 115ae65409dd..00bc662bcc49 100644 --- a/python/semantic_kernel/connectors/ai/bedrock/services/bedrock_base.py +++ b/python/semantic_kernel/connectors/ai/bedrock/services/bedrock_base.py @@ -4,6 +4,8 @@ from functools import partial from typing import Any, ClassVar +import boto3 + from semantic_kernel.connectors.ai.bedrock.services.model_provider.utils import run_in_executor from semantic_kernel.kernel_pydantic import KernelBaseModel @@ -19,6 +21,26 @@ class BedrockBase(KernelBaseModel, ABC): # Client: Use for model management bedrock_client: Any + def __init__( + self, + *, + runtime_client: Any | None = None, + client: Any | None = None, + **kwargs: Any, + ) -> None: + """Initialize the Amazon Bedrock Base Class. + + Args: + runtime_client: The Amazon Bedrock runtime client to use. + client: The Amazon Bedrock client to use. + **kwargs: Additional keyword arguments. + """ + super().__init__( + bedrock_runtime_client=runtime_client or boto3.client("bedrock-runtime"), + bedrock_client=client or boto3.client("bedrock"), + **kwargs, + ) + async def get_foundation_model_info(self, model_id: str) -> dict[str, Any]: """Get the foundation model information.""" response = await run_in_executor( diff --git a/python/semantic_kernel/connectors/ai/bedrock/services/bedrock_chat_completion.py b/python/semantic_kernel/connectors/ai/bedrock/services/bedrock_chat_completion.py index 5c4f3e6cd192..6768c603cd46 100644 --- a/python/semantic_kernel/connectors/ai/bedrock/services/bedrock_chat_completion.py +++ b/python/semantic_kernel/connectors/ai/bedrock/services/bedrock_chat_completion.py @@ -5,8 +5,6 @@ from functools import partial from typing import TYPE_CHECKING, Any, ClassVar -import boto3 - if sys.version_info >= (3, 12): from typing import override # pragma: no cover else: @@ -95,8 +93,8 @@ def __init__( super().__init__( ai_model_id=bedrock_settings.chat_model_id, service_id=service_id or bedrock_settings.chat_model_id, - bedrock_runtime_client=runtime_client or boto3.client("bedrock-runtime"), - bedrock_client=client or boto3.client("bedrock"), + runtime_client=runtime_client, + client=client, ) # region Overriding base class methods diff --git a/python/semantic_kernel/connectors/ai/bedrock/services/bedrock_text_completion.py b/python/semantic_kernel/connectors/ai/bedrock/services/bedrock_text_completion.py index 81092a7c7fa4..84fd5a8c8dcb 100644 --- a/python/semantic_kernel/connectors/ai/bedrock/services/bedrock_text_completion.py +++ b/python/semantic_kernel/connectors/ai/bedrock/services/bedrock_text_completion.py @@ -6,7 +6,6 @@ from functools import partial from typing import TYPE_CHECKING, Any -import boto3 from pydantic import ValidationError if sys.version_info >= (3, 12): @@ -73,8 +72,8 @@ def __init__( super().__init__( ai_model_id=bedrock_settings.text_model_id, service_id=service_id or bedrock_settings.text_model_id, - bedrock_runtime_client=runtime_client or boto3.client("bedrock-runtime"), - bedrock_client=client or boto3.client("bedrock"), + runtime_client=runtime_client, + client=client, ) # region Overriding base class methods diff --git a/python/semantic_kernel/connectors/ai/bedrock/services/bedrock_text_embedding.py b/python/semantic_kernel/connectors/ai/bedrock/services/bedrock_text_embedding.py index f963db5c5f0b..b5245fde8b9d 100644 --- a/python/semantic_kernel/connectors/ai/bedrock/services/bedrock_text_embedding.py +++ b/python/semantic_kernel/connectors/ai/bedrock/services/bedrock_text_embedding.py @@ -6,7 +6,6 @@ from functools import partial from typing import TYPE_CHECKING, Any -import boto3 from numpy import array, ndarray from pydantic import ValidationError @@ -70,8 +69,8 @@ def __init__( super().__init__( ai_model_id=bedrock_settings.embedding_model_id, service_id=service_id or bedrock_settings.embedding_model_id, - bedrock_runtime_client=runtime_client or boto3.client("bedrock-runtime"), - bedrock_client=client or boto3.client("bedrock"), + runtime_client=runtime_client, + client=client, ) @override From fa7d8a54fd70a9b0777ec87667ef488613257ed0 Mon Sep 17 00:00:00 2001 From: Tao Chen Date: Mon, 3 Mar 2025 14:30:00 -0800 Subject: [PATCH 2/6] Improve Bedrock agent setup docs --- .../concepts/agents/bedrock_agent/README.md | 32 +++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/python/samples/concepts/agents/bedrock_agent/README.md b/python/samples/concepts/agents/bedrock_agent/README.md index ea731c4e29c8..a07423a0e2d7 100644 --- a/python/samples/concepts/agents/bedrock_agent/README.md +++ b/python/samples/concepts/agents/bedrock_agent/README.md @@ -5,6 +5,38 @@ 1. You need to have an AWS account and [access to the foundation models](https://docs.aws.amazon.com/bedrock/latest/userguide/model-access-permissions.html) 2. [AWS CLI installed](https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html) and [configured](https://boto3.amazonaws.com/v1/documentation/api/latest/guide/quickstart.html#configuration) +### Configuration + +Follow this [guide](https://boto3.amazonaws.com/v1/documentation/api/latest/guide/quickstart.html#configuration) to configure your environment to use the Bedrock API. + +> Please configure the `aws_access_key_id`, `aws_secret_access_key`, and `region` otherwise you will need to create custom clients for the services. For example: + + ```python + runtime_client=boto.client( + "bedrock-runtime", + aws_access_key_id="your_access_key", + aws_secret_access_key="your_secret_key", + region_name="your_region", + [...other parameters you may need...] + ) + client=boto.client( + "bedrock", + aws_access_key_id="your_access_key", + aws_secret_access_key="your_secret_key", + region_name="your_region", + [...other parameters you may need...] + ) + + bedrock_agent = BedrockAgent.create_and_prepare_agent( + name="your_agent_name", + instructions="your_instructions", + runtime_client=runtime_client, + client=client, + [...other parameters you may need...] + ) + ``` + + ## Samples | Sample | Description | From a62bac1e2aaa5122c371ee96127b9592d51eb7ee Mon Sep 17 00:00:00 2001 From: Tao Chen Date: Mon, 3 Mar 2025 14:30:18 -0800 Subject: [PATCH 3/6] Remove empty line --- python/samples/concepts/agents/bedrock_agent/README.md | 1 - 1 file changed, 1 deletion(-) diff --git a/python/samples/concepts/agents/bedrock_agent/README.md b/python/samples/concepts/agents/bedrock_agent/README.md index a07423a0e2d7..c08eaab3a9bb 100644 --- a/python/samples/concepts/agents/bedrock_agent/README.md +++ b/python/samples/concepts/agents/bedrock_agent/README.md @@ -36,7 +36,6 @@ Follow this [guide](https://boto3.amazonaws.com/v1/documentation/api/latest/guid ) ``` - ## Samples | Sample | Description | From a07278a67b4de9fd0dfc86fbe4c985811b09efa3 Mon Sep 17 00:00:00 2001 From: Tao Chen Date: Mon, 3 Mar 2025 14:34:37 -0800 Subject: [PATCH 4/6] Reformat readme --- .../concepts/agents/bedrock_agent/README.md | 52 +++++++++---------- .../connectors/ai/bedrock/README.md | 40 +++++++------- 2 files changed, 46 insertions(+), 46 deletions(-) diff --git a/python/samples/concepts/agents/bedrock_agent/README.md b/python/samples/concepts/agents/bedrock_agent/README.md index c08eaab3a9bb..3e72751eb308 100644 --- a/python/samples/concepts/agents/bedrock_agent/README.md +++ b/python/samples/concepts/agents/bedrock_agent/README.md @@ -9,32 +9,32 @@ Follow this [guide](https://boto3.amazonaws.com/v1/documentation/api/latest/guide/quickstart.html#configuration) to configure your environment to use the Bedrock API. -> Please configure the `aws_access_key_id`, `aws_secret_access_key`, and `region` otherwise you will need to create custom clients for the services. For example: - - ```python - runtime_client=boto.client( - "bedrock-runtime", - aws_access_key_id="your_access_key", - aws_secret_access_key="your_secret_key", - region_name="your_region", - [...other parameters you may need...] - ) - client=boto.client( - "bedrock", - aws_access_key_id="your_access_key", - aws_secret_access_key="your_secret_key", - region_name="your_region", - [...other parameters you may need...] - ) - - bedrock_agent = BedrockAgent.create_and_prepare_agent( - name="your_agent_name", - instructions="your_instructions", - runtime_client=runtime_client, - client=client, - [...other parameters you may need...] - ) - ``` +Please configure the `aws_access_key_id`, `aws_secret_access_key`, and `region` otherwise you will need to create custom clients for the services. For example: + +```python +runtime_client=boto.client( + "bedrock-runtime", + aws_access_key_id="your_access_key", + aws_secret_access_key="your_secret_key", + region_name="your_region", + [...other parameters you may need...] +) +client=boto.client( + "bedrock", + aws_access_key_id="your_access_key", + aws_secret_access_key="your_secret_key", + region_name="your_region", + [...other parameters you may need...] +) + +bedrock_agent = BedrockAgent.create_and_prepare_agent( + name="your_agent_name", + instructions="your_instructions", + runtime_client=runtime_client, + client=client, + [...other parameters you may need...] +) +``` ## Samples diff --git a/python/semantic_kernel/connectors/ai/bedrock/README.md b/python/semantic_kernel/connectors/ai/bedrock/README.md index 6dee8f9c5bbb..bc9461aef7e8 100644 --- a/python/semantic_kernel/connectors/ai/bedrock/README.md +++ b/python/semantic_kernel/connectors/ai/bedrock/README.md @@ -11,26 +11,26 @@ Follow this [guide](https://boto3.amazonaws.com/v1/documentation/api/latest/guide/quickstart.html#configuration) to configure your environment to use the Bedrock API. -> Please configure the `aws_access_key_id`, `aws_secret_access_key`, and `region` otherwise you will need to create custom clients for the services. For example: - - ```python - runtime_client=boto.client( - "bedrock-runtime", - aws_access_key_id="your_access_key", - aws_secret_access_key="your_secret_key", - region_name="your_region", - [...other parameters you may need...] - ) - client=boto.client( - "bedrock", - aws_access_key_id="your_access_key", - aws_secret_access_key="your_secret_key", - region_name="your_region", - [...other parameters you may need...] - ) - - bedrock_chat_completion_service = BedrockChatCompletion(runtime_client=runtime_client, client=client) - ``` +Please configure the `aws_access_key_id`, `aws_secret_access_key`, and `region` otherwise you will need to create custom clients for the services. For example: + +```python +runtime_client=boto.client( + "bedrock-runtime", + aws_access_key_id="your_access_key", + aws_secret_access_key="your_secret_key", + region_name="your_region", + [...other parameters you may need...] +) +client=boto.client( + "bedrock", + aws_access_key_id="your_access_key", + aws_secret_access_key="your_secret_key", + region_name="your_region", + [...other parameters you may need...] +) + +bedrock_chat_completion_service = BedrockChatCompletion(runtime_client=runtime_client, client=client) +``` ## Supports From 6e26c7ed2da7a5f42b3a758e34ce37bf3e5616a1 Mon Sep 17 00:00:00 2001 From: Tao Chen Date: Mon, 3 Mar 2025 15:18:12 -0800 Subject: [PATCH 5/6] Reenable bedrock integration tests --- .../completions/chat_completion_test_base.py | 26 ++++++++++++------- .../completions/test_chat_completions.py | 2 -- .../completions/test_text_completion.py | 24 +++++++++++------ 3 files changed, 33 insertions(+), 19 deletions(-) diff --git a/python/tests/integration/completions/chat_completion_test_base.py b/python/tests/integration/completions/chat_completion_test_base.py index b5bf42111283..66f819798e8b 100644 --- a/python/tests/integration/completions/chat_completion_test_base.py +++ b/python/tests/integration/completions/chat_completion_test_base.py @@ -67,8 +67,6 @@ ["ONNX_GEN_AI_CHAT_MODEL_FOLDER"], raise_if_not_set=False ) # Tests are optional for ONNX anthropic_setup: bool = is_service_setup_for_testing(["ANTHROPIC_API_KEY", "ANTHROPIC_CHAT_MODEL_ID"]) -# When testing Bedrock, after logging into AWS CLI this has been set, so we can use it to check if the service is setup -bedrock_setup: bool = is_service_setup_for_testing(["AWS_DEFAULT_REGION"], raise_if_not_set=False) # A mock plugin that contains a function that returns a complex object. @@ -90,7 +88,7 @@ class ChatCompletionTestBase(CompletionTestBase): """Base class for testing completion services.""" @override - @pytest.fixture(scope="function") + @pytest.fixture(scope="class") def services(self) -> dict[str, tuple[ServiceType | None, type[PromptExecutionSettings] | None]]: azure_openai_setup = True azure_openai_settings = AzureOpenAISettings.create() @@ -152,27 +150,27 @@ def services(self) -> dict[str, tuple[ServiceType | None, type[PromptExecutionSe OnnxGenAIPromptExecutionSettings, ), "bedrock_amazon_titan": ( - BedrockChatCompletion(model_id="amazon.titan-text-premier-v1:0") if bedrock_setup else None, + self._try_create_bedrock_chat_completion_client("amazon.titan-text-premier-v1:0"), BedrockChatPromptExecutionSettings, ), "bedrock_ai21labs": ( - BedrockChatCompletion(model_id="ai21.jamba-1-5-mini-v1:0") if bedrock_setup else None, + self._try_create_bedrock_chat_completion_client("ai21.jamba-1-5-mini-v1:0"), BedrockChatPromptExecutionSettings, ), "bedrock_anthropic_claude": ( - BedrockChatCompletion(model_id="anthropic.claude-3-5-sonnet-20240620-v1:0") if bedrock_setup else None, + self._try_create_bedrock_chat_completion_client("anthropic.claude-3-5-sonnet-20240620-v1:0"), BedrockChatPromptExecutionSettings, ), "bedrock_cohere_command": ( - BedrockChatCompletion(model_id="cohere.command-r-v1:0") if bedrock_setup else None, + self._try_create_bedrock_chat_completion_client("cohere.command-r-v1:0"), BedrockChatPromptExecutionSettings, ), "bedrock_meta_llama": ( - BedrockChatCompletion(model_id="meta.llama3-70b-instruct-v1:0") if bedrock_setup else None, + self._try_create_bedrock_chat_completion_client("meta.llama3-70b-instruct-v1:0"), BedrockChatPromptExecutionSettings, ), "bedrock_mistralai": ( - BedrockChatCompletion(model_id="mistral.mistral-small-2402-v1:0") if bedrock_setup else None, + self._try_create_bedrock_chat_completion_client("mistral.mistral-small-2402-v1:0"), BedrockChatPromptExecutionSettings, ), } @@ -218,3 +216,13 @@ async def get_chat_completion_response( if parts: return sum(parts[1:], parts[0]) raise AssertionError("No response") + + def _try_create_bedrock_chat_completion_client(self, model_id: str) -> BedrockChatCompletion | None: + try: + return BedrockChatCompletion(model_id=model_id) + except Exception as ex: + from conftest import logger + + logger.warning(ex) + # Returning None so that the test that uses this service will be skipped + return None diff --git a/python/tests/integration/completions/test_chat_completions.py b/python/tests/integration/completions/test_chat_completions.py index 937584957bff..f0e9a4fdb5ea 100644 --- a/python/tests/integration/completions/test_chat_completions.py +++ b/python/tests/integration/completions/test_chat_completions.py @@ -18,7 +18,6 @@ from tests.integration.completions.chat_completion_test_base import ( ChatCompletionTestBase, anthropic_setup, - bedrock_setup, mistral_ai_setup, ollama_setup, onnx_setup, @@ -190,7 +189,6 @@ class Reasoning(KernelBaseModel): ChatMessageContent(role=AuthorRole.USER, items=[TextContent(text="How are you today?")]), ], {}, - marks=pytest.mark.skipif(not bedrock_setup, reason="Bedrock Environment Variables not set"), id="bedrock_amazon_titan_text_input", ), pytest.param( diff --git a/python/tests/integration/completions/test_text_completion.py b/python/tests/integration/completions/test_text_completion.py index ca0844cee59f..c2e1e06d523a 100644 --- a/python/tests/integration/completions/test_text_completion.py +++ b/python/tests/integration/completions/test_text_completion.py @@ -45,7 +45,6 @@ onnx_setup: bool = is_service_setup_for_testing( ["ONNX_GEN_AI_TEXT_MODEL_FOLDER"], raise_if_not_set=False ) # Tests are optional for ONNX -bedrock_setup = is_service_setup_for_testing(["AWS_DEFAULT_REGION"], raise_if_not_set=False) pytestmark = pytest.mark.parametrize( "service_id, execution_settings_kwargs, inputs, kwargs", @@ -138,7 +137,6 @@ {}, ["Repeat the word Hello once"], {}, - marks=pytest.mark.skipif(not bedrock_setup, reason="Not setup"), id="bedrock_amazon_titan_text_completion", ), pytest.param( @@ -255,27 +253,27 @@ def services(self) -> dict[str, tuple[ServiceType | None, type[PromptExecutionSe # Amazon Bedrock supports models from multiple providers but requests to and responses from the models are # inconsistent. So we need to test each model separately. "bedrock_amazon_titan": ( - BedrockTextCompletion(model_id="amazon.titan-text-premier-v1:0") if bedrock_setup else None, + self._try_create_bedrock_text_completion_client("amazon.titan-text-premier-v1:0"), BedrockTextPromptExecutionSettings, ), "bedrock_anthropic_claude": ( - BedrockTextCompletion(model_id="anthropic.claude-v2") if bedrock_setup else None, + self._try_create_bedrock_text_completion_client("anthropic.claude-v2"), BedrockTextPromptExecutionSettings, ), "bedrock_cohere_command": ( - BedrockTextCompletion(model_id="cohere.command-text-v14") if bedrock_setup else None, + self._try_create_bedrock_text_completion_client("cohere.command-text-v14"), BedrockTextPromptExecutionSettings, ), "bedrock_ai21labs": ( - BedrockTextCompletion(model_id="ai21.j2-mid-v1") if bedrock_setup else None, + self._try_create_bedrock_text_completion_client("ai21.j2-mid-v1"), BedrockTextPromptExecutionSettings, ), "bedrock_meta_llama": ( - BedrockTextCompletion(model_id="meta.llama3-70b-instruct-v1:0") if bedrock_setup else None, + self._try_create_bedrock_text_completion_client("meta.llama3-70b-instruct-v1:0"), BedrockTextPromptExecutionSettings, ), "bedrock_mistralai": ( - BedrockTextCompletion(model_id="mistral.mistral-7b-instruct-v0:2") if bedrock_setup else None, + self._try_create_bedrock_text_completion_client("mistral.mistral-7b-instruct-v0:2"), BedrockTextPromptExecutionSettings, ), } @@ -373,3 +371,13 @@ async def _test_helper( name="text completions", ) self.evaluate(response) + + def _try_create_bedrock_text_completion_client(self, model_id: str) -> BedrockTextCompletion | None: + try: + return BedrockTextCompletion(model_id=model_id) + except Exception as ex: + from conftest import logger + + logger.warning(ex) + # Returning None so that the test that uses this service will be skipped + return None From f3a5e66726f0c38f47526d04d5ea228c467701d5 Mon Sep 17 00:00:00 2001 From: Tao Chen Date: Mon, 3 Mar 2025 16:17:50 -0800 Subject: [PATCH 6/6] Fix integration tests --- .../integration/completions/chat_completion_test_base.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/tests/integration/completions/chat_completion_test_base.py b/python/tests/integration/completions/chat_completion_test_base.py index 66f819798e8b..890298556605 100644 --- a/python/tests/integration/completions/chat_completion_test_base.py +++ b/python/tests/integration/completions/chat_completion_test_base.py @@ -88,7 +88,9 @@ class ChatCompletionTestBase(CompletionTestBase): """Base class for testing completion services.""" @override - @pytest.fixture(scope="class") + @pytest.fixture( + scope="function" + ) # This needs to be scoped to function to avoid resources getting cleaned up after each test def services(self) -> dict[str, tuple[ServiceType | None, type[PromptExecutionSettings] | None]]: azure_openai_setup = True azure_openai_settings = AzureOpenAISettings.create()