chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,371 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
|
||||
from functools import reduce
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import boto3
|
||||
import pytest
|
||||
|
||||
from semantic_kernel.connectors.ai.bedrock.bedrock_prompt_execution_settings import BedrockChatPromptExecutionSettings
|
||||
from semantic_kernel.connectors.ai.bedrock.services.bedrock_chat_completion import BedrockChatCompletion
|
||||
from semantic_kernel.connectors.ai.bedrock.services.model_provider.bedrock_model_provider import BedrockModelProvider
|
||||
from semantic_kernel.connectors.ai.completion_usage import CompletionUsage
|
||||
from semantic_kernel.contents.chat_history import ChatHistory
|
||||
from semantic_kernel.contents.chat_message_content import ChatMessageContent
|
||||
from semantic_kernel.contents.streaming_chat_message_content import StreamingChatMessageContent
|
||||
from semantic_kernel.contents.text_content import TextContent
|
||||
from semantic_kernel.contents.utils.author_role import AuthorRole
|
||||
from semantic_kernel.contents.utils.finish_reason import FinishReason
|
||||
from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError, ServiceInvalidResponseError
|
||||
from tests.unit.connectors.ai.bedrock.conftest import MockBedrockClient, MockBedrockRuntimeClient
|
||||
|
||||
# region init
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_bedrock_chat_completion_init(mock_client, bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Chat Completion service"""
|
||||
bedrock_chat_completion = BedrockChatCompletion()
|
||||
|
||||
assert bedrock_chat_completion.ai_model_id == bedrock_unit_test_env["BEDROCK_CHAT_MODEL_ID"]
|
||||
assert bedrock_chat_completion.service_id == bedrock_unit_test_env["BEDROCK_CHAT_MODEL_ID"]
|
||||
|
||||
assert bedrock_chat_completion.bedrock_model_provider == BedrockModelProvider(
|
||||
bedrock_unit_test_env["BEDROCK_MODEL_PROVIDER"]
|
||||
)
|
||||
assert bedrock_chat_completion.bedrock_client is not None
|
||||
assert bedrock_chat_completion.bedrock_runtime_client is not None
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_bedrock_chat_completion_init_model_id_override(mock_client, bedrock_unit_test_env, model_id) -> None:
|
||||
"""Test initialization of Amazon Bedrock Chat Completion service"""
|
||||
bedrock_chat_completion = BedrockChatCompletion(model_id=model_id)
|
||||
|
||||
assert bedrock_chat_completion.ai_model_id == model_id
|
||||
assert bedrock_chat_completion.service_id == model_id
|
||||
|
||||
assert bedrock_chat_completion.bedrock_client is not None
|
||||
assert bedrock_chat_completion.bedrock_runtime_client is not None
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_bedrock_chat_completion_init_custom_service_id(mock_client, bedrock_unit_test_env, service_id) -> None:
|
||||
"""Test initialization of Amazon Bedrock Chat Completion service"""
|
||||
bedrock_chat_completion = BedrockChatCompletion(service_id=service_id)
|
||||
|
||||
assert bedrock_chat_completion.service_id == service_id
|
||||
|
||||
assert bedrock_chat_completion.bedrock_client is not None
|
||||
assert bedrock_chat_completion.bedrock_runtime_client is not None
|
||||
|
||||
|
||||
def test_bedrock_chat_completion_init_custom_clients(bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Chat Completion service"""
|
||||
bedrock_chat_completion = BedrockChatCompletion(
|
||||
runtime_client=MockBedrockRuntimeClient(),
|
||||
client=MockBedrockClient(),
|
||||
)
|
||||
|
||||
assert isinstance(bedrock_chat_completion.bedrock_client, MockBedrockClient)
|
||||
assert isinstance(bedrock_chat_completion.bedrock_runtime_client, MockBedrockRuntimeClient)
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_bedrock_chat_completion_init_custom_client(mock_client, bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Chat Completion service"""
|
||||
bedrock_chat_completion = BedrockChatCompletion(
|
||||
client=MockBedrockClient(),
|
||||
)
|
||||
|
||||
assert isinstance(bedrock_chat_completion.bedrock_client, MockBedrockClient)
|
||||
assert bedrock_chat_completion.bedrock_runtime_client is not None
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_bedrock_chat_completion_init_custom_runtime_client(mock_client, bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Chat Completion service"""
|
||||
bedrock_chat_completion = BedrockChatCompletion(
|
||||
runtime_client=MockBedrockRuntimeClient(),
|
||||
)
|
||||
|
||||
assert bedrock_chat_completion.bedrock_client is not None
|
||||
assert isinstance(bedrock_chat_completion.bedrock_runtime_client, MockBedrockRuntimeClient)
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_bedrock_chat_completion_init_custom_bedrock_model_provider(mock_client, bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Chat Completion service"""
|
||||
bedrock_chat_completion = BedrockChatCompletion(
|
||||
model_provider=BedrockModelProvider.AMAZON,
|
||||
)
|
||||
|
||||
assert bedrock_chat_completion.bedrock_model_provider == BedrockModelProvider.AMAZON
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["BEDROCK_CHAT_MODEL_ID"]], indirect=True)
|
||||
def test_bedrock_chat_completion_client_init_with_empty_model_id(bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Chat Completion service with empty model id"""
|
||||
with pytest.raises(ServiceInitializationError, match="The Amazon Bedrock Chat Model ID is missing."):
|
||||
BedrockChatCompletion(env_file_path="fake_env_file_path.env")
|
||||
|
||||
|
||||
def test_bedrock_chat_completion_client_init_invalid_settings(bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Chat Completion service with invalid settings"""
|
||||
with pytest.raises(
|
||||
ServiceInitializationError, match="Failed to initialize the Amazon Bedrock Chat Completion Service."
|
||||
):
|
||||
BedrockChatCompletion(model_id=123) # Model ID must be a string
|
||||
|
||||
|
||||
def test_bedrock_chat_completion_client_init_invalid_model_provider(bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Chat Completion service with invalid settings"""
|
||||
with pytest.raises(
|
||||
ServiceInitializationError, match="Failed to initialize the Amazon Bedrock Chat Completion Service."
|
||||
):
|
||||
BedrockChatCompletion(model_provider="invalid_provider")
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_prompt_execution_settings_class(mock_client, bedrock_unit_test_env) -> None:
|
||||
"""Test getting prompt execution settings class"""
|
||||
bedrock_completion_client = BedrockChatCompletion()
|
||||
assert bedrock_completion_client.get_prompt_execution_settings_class() == BedrockChatPromptExecutionSettings
|
||||
|
||||
|
||||
# endregion
|
||||
|
||||
# region private methods
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_prepare_chat_history_for_request(mock_client, bedrock_unit_test_env, chat_history) -> None:
|
||||
"""Test preparing chat history for request"""
|
||||
bedrock_chat_completion = BedrockChatCompletion()
|
||||
parsed_chat_history = bedrock_chat_completion._prepare_chat_history_for_request(chat_history)
|
||||
|
||||
assert isinstance(parsed_chat_history, list)
|
||||
assert len(parsed_chat_history) == len(chat_history) - 2 # Exclude system message
|
||||
assert all([item["role"] in ["user", "assistant"] for item in parsed_chat_history])
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_prepare_system_message_for_request(mock_client, bedrock_unit_test_env, chat_history) -> None:
|
||||
"""Test preparing system message for request"""
|
||||
bedrock_chat_completion = BedrockChatCompletion()
|
||||
parsed_system_message = bedrock_chat_completion._prepare_system_messages_for_request(chat_history)
|
||||
|
||||
assert isinstance(parsed_system_message, list)
|
||||
assert len(parsed_system_message) == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_id",
|
||||
[
|
||||
"amazon.titan",
|
||||
"anthropic.claude",
|
||||
"cohere.command",
|
||||
"ai21.jamba",
|
||||
"meta.llama",
|
||||
"mistral.ai",
|
||||
],
|
||||
)
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_prepare_settings_for_request(mock_client, model_id, chat_history) -> None:
|
||||
"""Test preparing settings for request"""
|
||||
bedrock_chat_completion = BedrockChatCompletion(model_id=model_id)
|
||||
settings = BedrockChatPromptExecutionSettings()
|
||||
parsed_settings = bedrock_chat_completion._prepare_settings_for_request(chat_history, settings)
|
||||
|
||||
assert isinstance(parsed_settings, dict)
|
||||
assert parsed_settings["modelId"] == bedrock_chat_completion.ai_model_id
|
||||
assert parsed_settings["messages"] == bedrock_chat_completion._prepare_chat_history_for_request(chat_history)
|
||||
assert parsed_settings["system"] == bedrock_chat_completion._prepare_system_messages_for_request(chat_history)
|
||||
assert isinstance(parsed_settings["inferenceConfig"], dict)
|
||||
assert all([parsed_settings["inferenceConfig"].values()])
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_id",
|
||||
[
|
||||
"arn:aws:bedrock:us-east-1:972143716085:application-inference-profile/123456",
|
||||
],
|
||||
)
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_prepare_settings_for_request_with_application_inference_profile(mock_client, model_id, chat_history) -> None:
|
||||
"""Test preparing settings for request"""
|
||||
# Without a valid model provider, it should raise an error
|
||||
bedrock_chat_completion = BedrockChatCompletion(model_id=model_id)
|
||||
settings = BedrockChatPromptExecutionSettings()
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=f"Model ID {model_id} does not contain a valid model provider name.",
|
||||
):
|
||||
bedrock_chat_completion._prepare_settings_for_request(chat_history, settings)
|
||||
|
||||
# With a valid model provider, it should not raise an error
|
||||
bedrock_chat_completion = BedrockChatCompletion(model_id=model_id, model_provider=BedrockModelProvider.AMAZON)
|
||||
parsed_settings = bedrock_chat_completion._prepare_settings_for_request(chat_history, settings)
|
||||
|
||||
assert isinstance(parsed_settings, dict)
|
||||
assert parsed_settings["modelId"] == bedrock_chat_completion.ai_model_id
|
||||
assert parsed_settings["messages"] == bedrock_chat_completion._prepare_chat_history_for_request(chat_history)
|
||||
assert parsed_settings["system"] == bedrock_chat_completion._prepare_system_messages_for_request(chat_history)
|
||||
assert isinstance(parsed_settings["inferenceConfig"], dict)
|
||||
assert all([parsed_settings["inferenceConfig"].values()])
|
||||
|
||||
|
||||
# endregion
|
||||
|
||||
|
||||
# region chat completion
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
# These are fake model ids with the supported prefixes
|
||||
"model_id",
|
||||
[
|
||||
"amazon.titan",
|
||||
"anthropic.claude",
|
||||
"cohere.command",
|
||||
"ai21.jamba",
|
||||
"meta.llama",
|
||||
"mistral.ai",
|
||||
],
|
||||
)
|
||||
async def test_bedrock_chat_completion(
|
||||
model_id,
|
||||
chat_history: ChatHistory,
|
||||
mock_bedrock_chat_completion_response,
|
||||
) -> None:
|
||||
"""Test Amazon Bedrock Chat Completion complete method"""
|
||||
with patch.object(
|
||||
MockBedrockRuntimeClient, "converse", return_value=mock_bedrock_chat_completion_response
|
||||
) as mock_converse:
|
||||
# Setup
|
||||
bedrock_chat_completion = BedrockChatCompletion(
|
||||
model_id=model_id,
|
||||
runtime_client=MockBedrockRuntimeClient(),
|
||||
client=MockBedrockClient(),
|
||||
)
|
||||
|
||||
# Act
|
||||
settings = BedrockChatPromptExecutionSettings()
|
||||
response = await bedrock_chat_completion.get_chat_message_contents(chat_history=chat_history, settings=settings)
|
||||
|
||||
# Assert
|
||||
mock_converse.assert_called_once_with(
|
||||
**(bedrock_chat_completion._prepare_settings_for_request(chat_history, settings))
|
||||
)
|
||||
|
||||
assert isinstance(response, list)
|
||||
assert len(response) == 1
|
||||
assert isinstance(response[0], ChatMessageContent)
|
||||
assert response[0].ai_model_id == model_id
|
||||
assert response[0].role == AuthorRole.ASSISTANT
|
||||
assert len(response[0].items) == 1
|
||||
assert isinstance(response[0].items[0], TextContent)
|
||||
assert response[0].finish_reason == FinishReason.STOP
|
||||
assert response[0].metadata["usage"] == CompletionUsage(
|
||||
prompt_tokens=mock_bedrock_chat_completion_response["usage"]["inputTokens"],
|
||||
completion_tokens=mock_bedrock_chat_completion_response["usage"]["outputTokens"],
|
||||
)
|
||||
assert (
|
||||
response[0].items[0].text
|
||||
== mock_bedrock_chat_completion_response["output"]["message"]["content"][0]["text"]
|
||||
)
|
||||
assert response[0].inner_content == mock_bedrock_chat_completion_response
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
# These are fake model ids with the supported prefixes
|
||||
"model_id",
|
||||
[
|
||||
"amazon.titan",
|
||||
"anthropic.claude",
|
||||
"cohere.command",
|
||||
"ai21.jamba",
|
||||
"meta.llama",
|
||||
"mistral.ai",
|
||||
],
|
||||
)
|
||||
async def test_bedrock_streaming_chat_completion(
|
||||
model_id,
|
||||
chat_history: ChatHistory,
|
||||
mock_bedrock_streaming_chat_completion_response,
|
||||
) -> None:
|
||||
"""Test Amazon Bedrock Streaming Chat Completion complete method"""
|
||||
with patch.object(
|
||||
MockBedrockRuntimeClient, "converse_stream", return_value=mock_bedrock_streaming_chat_completion_response
|
||||
) as mock_converse_stream:
|
||||
# Setup
|
||||
bedrock_chat_completion = BedrockChatCompletion(
|
||||
model_id=model_id,
|
||||
runtime_client=MockBedrockRuntimeClient(),
|
||||
client=MockBedrockClient(),
|
||||
)
|
||||
|
||||
# Act
|
||||
settings = BedrockChatPromptExecutionSettings()
|
||||
chunks: list[StreamingChatMessageContent] = []
|
||||
async for streaming_messages in bedrock_chat_completion.get_streaming_chat_message_contents(
|
||||
chat_history=chat_history, settings=settings
|
||||
):
|
||||
chunks.extend(streaming_messages)
|
||||
response = reduce(lambda p, r: p + r, chunks)
|
||||
|
||||
# Assert
|
||||
mock_converse_stream.assert_called_once_with(
|
||||
**(bedrock_chat_completion._prepare_settings_for_request(chat_history, settings))
|
||||
)
|
||||
|
||||
assert isinstance(response, StreamingChatMessageContent)
|
||||
assert response.ai_model_id == model_id
|
||||
assert response.role == AuthorRole.ASSISTANT
|
||||
assert len(response.items) == 1
|
||||
assert isinstance(response.inner_content, list)
|
||||
assert len(response.inner_content) == 7
|
||||
assert response.finish_reason == FinishReason.STOP
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
# These are fake model ids with the supported prefixes
|
||||
"model_id",
|
||||
[
|
||||
"amazon.titan",
|
||||
"anthropic.claude",
|
||||
"cohere.command",
|
||||
"ai21.jamba",
|
||||
"meta.llama",
|
||||
"mistral.ai",
|
||||
],
|
||||
)
|
||||
async def test_bedrock_streaming_chat_completion_invalid_event(
|
||||
model_id,
|
||||
chat_history: ChatHistory,
|
||||
mock_bedrock_streaming_chat_completion_invalid_response,
|
||||
) -> None:
|
||||
"""Test Amazon Bedrock Streaming Chat Completion complete method"""
|
||||
with patch.object(
|
||||
MockBedrockRuntimeClient,
|
||||
"converse_stream",
|
||||
return_value=mock_bedrock_streaming_chat_completion_invalid_response,
|
||||
):
|
||||
# Setup
|
||||
bedrock_chat_completion = BedrockChatCompletion(
|
||||
model_id=model_id,
|
||||
runtime_client=MockBedrockRuntimeClient(),
|
||||
client=MockBedrockClient(),
|
||||
)
|
||||
|
||||
# Act
|
||||
settings = BedrockChatPromptExecutionSettings()
|
||||
with pytest.raises(ServiceInvalidResponseError):
|
||||
async for chunk in bedrock_chat_completion.get_streaming_chat_message_contents(
|
||||
chat_history=chat_history, settings=settings
|
||||
):
|
||||
pass
|
||||
|
||||
|
||||
# endregion
|
||||
+443
@@ -0,0 +1,443 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from semantic_kernel.connectors.ai.bedrock.bedrock_prompt_execution_settings import BedrockChatPromptExecutionSettings
|
||||
from semantic_kernel.connectors.ai.bedrock.services.model_provider.bedrock_model_provider import (
|
||||
BedrockModelProvider,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.bedrock.services.model_provider.utils import (
|
||||
MESSAGE_CONVERTERS,
|
||||
finish_reason_from_bedrock_to_semantic_kernel,
|
||||
remove_none_recursively,
|
||||
update_settings_from_function_choice_configuration,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.function_call_choice_configuration import FunctionCallChoiceConfiguration
|
||||
from semantic_kernel.connectors.ai.function_choice_behavior import FunctionChoiceBehavior
|
||||
from semantic_kernel.connectors.ai.function_choice_type import FunctionChoiceType
|
||||
from semantic_kernel.contents.chat_message_content import ChatMessageContent
|
||||
from semantic_kernel.contents.function_call_content import FunctionCallContent
|
||||
from semantic_kernel.contents.function_result_content import FunctionResultContent
|
||||
from semantic_kernel.contents.image_content import ImageContent
|
||||
from semantic_kernel.contents.text_content import TextContent
|
||||
from semantic_kernel.contents.utils.author_role import AuthorRole
|
||||
from semantic_kernel.contents.utils.finish_reason import FinishReason
|
||||
from semantic_kernel.exceptions.service_exceptions import ServiceInvalidRequestError
|
||||
from semantic_kernel.kernel import Kernel
|
||||
|
||||
|
||||
def test_remove_none_recursively():
|
||||
data = {
|
||||
"a": 1,
|
||||
"b": None,
|
||||
"c": {
|
||||
"d": 2,
|
||||
"e": None,
|
||||
"f": {
|
||||
"g": 3,
|
||||
"h": None,
|
||||
},
|
||||
},
|
||||
}
|
||||
expected = {
|
||||
"a": 1,
|
||||
"c": {
|
||||
"d": 2,
|
||||
"f": {
|
||||
"g": 3,
|
||||
},
|
||||
},
|
||||
}
|
||||
assert remove_none_recursively(data) == expected
|
||||
|
||||
|
||||
def test_remove_recursively_max_depth():
|
||||
data = {
|
||||
"a": {"b": None},
|
||||
}
|
||||
|
||||
assert remove_none_recursively(data, max_depth=1) == data
|
||||
|
||||
|
||||
def test_update_settings_from_function_choice_configuration_auto(kernel: Kernel, custom_plugin_class) -> None:
|
||||
kernel.add_plugin(plugin=custom_plugin_class(), plugin_name="custom_plugin")
|
||||
|
||||
settings = BedrockChatPromptExecutionSettings()
|
||||
|
||||
auto_function_choice_behavior = FunctionChoiceBehavior.Auto()
|
||||
auto_function_choice_behavior.configure(
|
||||
kernel,
|
||||
update_settings_from_function_choice_configuration,
|
||||
settings,
|
||||
)
|
||||
|
||||
assert "auto" in settings.tool_choice
|
||||
assert len(settings.tools) == 1
|
||||
|
||||
|
||||
def test_update_settings_from_function_choice_configuration_auto_without_plugin(kernel: Kernel) -> None:
|
||||
settings = BedrockChatPromptExecutionSettings()
|
||||
|
||||
auto_function_choice_behavior = FunctionChoiceBehavior.Auto()
|
||||
auto_function_choice_behavior.configure(
|
||||
kernel,
|
||||
update_settings_from_function_choice_configuration,
|
||||
settings,
|
||||
)
|
||||
|
||||
assert settings.tool_choice is None
|
||||
assert settings.tools is None
|
||||
|
||||
|
||||
def test_update_settings_from_function_choice_configuration_none(kernel: Kernel) -> None:
|
||||
settings = BedrockChatPromptExecutionSettings()
|
||||
|
||||
auto_function_choice_behavior = FunctionChoiceBehavior.NoneInvoke()
|
||||
auto_function_choice_behavior.configure(
|
||||
kernel,
|
||||
update_settings_from_function_choice_configuration,
|
||||
settings,
|
||||
)
|
||||
|
||||
assert settings.tool_choice is None
|
||||
assert settings.tools is None
|
||||
|
||||
|
||||
def test_update_settings_from_function_choice_configuration_required_with_one_function(
|
||||
kernel: Kernel,
|
||||
custom_plugin_class,
|
||||
) -> None:
|
||||
kernel.add_plugin(plugin=custom_plugin_class(), plugin_name="custom_plugin")
|
||||
|
||||
settings = BedrockChatPromptExecutionSettings()
|
||||
|
||||
auto_function_choice_behavior = FunctionChoiceBehavior.Required()
|
||||
auto_function_choice_behavior.configure(
|
||||
kernel,
|
||||
update_settings_from_function_choice_configuration,
|
||||
settings,
|
||||
)
|
||||
|
||||
assert "tool" in settings.tool_choice
|
||||
assert len(settings.tools) == 1
|
||||
|
||||
|
||||
def test_update_settings_from_function_choice_configuration_required_with_more_than_one_functions(
|
||||
kernel: Kernel,
|
||||
custom_plugin_class,
|
||||
experimental_plugin_class,
|
||||
) -> None:
|
||||
kernel.add_plugin(plugin=custom_plugin_class(), plugin_name="custom_plugin")
|
||||
kernel.add_plugin(plugin=experimental_plugin_class(), plugin_name="experimental_plugin")
|
||||
|
||||
settings = BedrockChatPromptExecutionSettings()
|
||||
|
||||
auto_function_choice_behavior = FunctionChoiceBehavior.Required()
|
||||
auto_function_choice_behavior.configure(
|
||||
kernel,
|
||||
update_settings_from_function_choice_configuration,
|
||||
settings,
|
||||
)
|
||||
|
||||
assert "any" in settings.tool_choice
|
||||
assert len(settings.tools) == 2
|
||||
|
||||
|
||||
def test_inference_profile_with_bedrock_model() -> None:
|
||||
"""Test the BedrockModelProvider class returns the correct model for a given inference profile."""
|
||||
|
||||
us_amazon_inference_profile = "us.amazon.nova-lite-v1:0"
|
||||
assert BedrockModelProvider.to_model_provider(us_amazon_inference_profile) == BedrockModelProvider.AMAZON
|
||||
|
||||
us_anthropic_inference_profile = "us.anthropic.claude-3-sonnet-20240229-v1:0"
|
||||
assert BedrockModelProvider.to_model_provider(us_anthropic_inference_profile) == BedrockModelProvider.ANTHROPIC
|
||||
|
||||
eu_meta_inference_profile = "eu.meta.llama3-2-3b-instruct-v1:0"
|
||||
assert BedrockModelProvider.to_model_provider(eu_meta_inference_profile) == BedrockModelProvider.META
|
||||
|
||||
unknown_inference_profile = "unknown"
|
||||
with pytest.raises(ValueError, match="Model ID unknown does not contain a valid model provider name."):
|
||||
BedrockModelProvider.to_model_provider(unknown_inference_profile)
|
||||
|
||||
|
||||
def test_remove_none_recursively_empty_dict() -> None:
|
||||
"""Test that an empty dict returns an empty dict."""
|
||||
assert remove_none_recursively({}) == {}
|
||||
|
||||
|
||||
def test_remove_none_recursively_no_none() -> None:
|
||||
"""Test that a dict with no None values remains the same."""
|
||||
original = {"a": 1, "b": 2}
|
||||
result = remove_none_recursively(original)
|
||||
assert result == {"a": 1, "b": 2}
|
||||
|
||||
|
||||
def test_remove_none_recursively_with_none() -> None:
|
||||
"""Test that dict values of None are removed."""
|
||||
original = {"a": 1, "b": None, "c": {"d": None, "e": 3}}
|
||||
result = remove_none_recursively(original)
|
||||
# 'b' should be removed and 'd' inside nested dict should be removed
|
||||
assert result == {"a": 1, "c": {"e": 3}}
|
||||
|
||||
|
||||
def test_remove_none_recursively_max_depth() -> None:
|
||||
"""Test that the function respects max_depth."""
|
||||
original = {"a": {"b": {"c": None}}}
|
||||
# If max_depth=1, it won't go deep enough to remove 'c'.
|
||||
result = remove_none_recursively(original, max_depth=1)
|
||||
assert result == {"a": {"b": {"c": None}}}
|
||||
|
||||
# If max_depth=3, it should remove 'c'.
|
||||
result = remove_none_recursively(original, max_depth=3)
|
||||
assert result == {"a": {"b": {}}}
|
||||
|
||||
|
||||
def test_format_system_message() -> None:
|
||||
"""Test that system message is formatted correctly."""
|
||||
content = ChatMessageContent(role=AuthorRole.SYSTEM, content="System message")
|
||||
formatted = MESSAGE_CONVERTERS[AuthorRole.SYSTEM](content)
|
||||
assert formatted == {"text": "System message"}
|
||||
|
||||
|
||||
def test_format_user_message_text_only() -> None:
|
||||
"""Test user message with only text content."""
|
||||
text_item = TextContent(text="Hello!")
|
||||
user_message = ChatMessageContent(role=AuthorRole.USER, items=[text_item])
|
||||
|
||||
formatted = MESSAGE_CONVERTERS[AuthorRole.USER](user_message)
|
||||
assert formatted["role"] == "user"
|
||||
assert len(formatted["content"]) == 1
|
||||
assert formatted["content"][0] == {"text": "Hello!"}
|
||||
|
||||
|
||||
def test_format_user_message_image_only() -> None:
|
||||
"""Test user message with only image content."""
|
||||
img_item = ImageContent(data=b"abc", mime_type="image/png")
|
||||
user_message = ChatMessageContent(role=AuthorRole.USER, items=[img_item])
|
||||
|
||||
formatted = MESSAGE_CONVERTERS[AuthorRole.USER](user_message)
|
||||
assert formatted["role"] == "user"
|
||||
assert len(formatted["content"]) == 1
|
||||
image_section = formatted["content"][0].get("image")
|
||||
assert image_section["format"] == "png"
|
||||
assert image_section["source"]["bytes"] == b"abc"
|
||||
|
||||
|
||||
def test_format_user_message_unsupported_content() -> None:
|
||||
"""Test user message raises error with unsupported content type."""
|
||||
# We can simulate an unsupported content type by using FunctionCallContent.
|
||||
func_call_item = FunctionCallContent(id="123", function_name="test_function", arguments="{}")
|
||||
user_message = ChatMessageContent(role=AuthorRole.USER, items=[func_call_item])
|
||||
|
||||
with pytest.raises(ServiceInvalidRequestError) as exc:
|
||||
MESSAGE_CONVERTERS[AuthorRole.USER](user_message)
|
||||
|
||||
assert "Only text and image content are supported in a user message." in str(exc.value)
|
||||
|
||||
|
||||
def test_format_assistant_message_text_content() -> None:
|
||||
"""Test assistant message with text content."""
|
||||
text_item = TextContent(text="Assistant response")
|
||||
assistant_message = ChatMessageContent(role=AuthorRole.ASSISTANT, items=[text_item])
|
||||
|
||||
formatted = MESSAGE_CONVERTERS[AuthorRole.ASSISTANT](assistant_message)
|
||||
assert formatted["role"] == "assistant"
|
||||
assert formatted["content"] == [{"text": "Assistant response"}]
|
||||
|
||||
|
||||
def test_format_assistant_message_function_call_content() -> None:
|
||||
"""Test assistant message with function call content."""
|
||||
func_item = FunctionCallContent(
|
||||
id="fc1", plugin_name="plugin", function_name="function", arguments='{"param": "value"}'
|
||||
)
|
||||
assistant_message = ChatMessageContent(role=AuthorRole.ASSISTANT, items=[func_item])
|
||||
|
||||
formatted = MESSAGE_CONVERTERS[AuthorRole.ASSISTANT](assistant_message)
|
||||
assert len(formatted["content"]) == 1
|
||||
tool_use = formatted["content"][0].get("toolUse")
|
||||
assert tool_use
|
||||
assert tool_use["toolUseId"] == "fc1"
|
||||
assert tool_use["name"] == "plugin-function"
|
||||
assert tool_use["input"] == {"param": "value"}
|
||||
|
||||
|
||||
def test_format_assistant_message_image_content_raises() -> None:
|
||||
"""Test assistant message with image raises error."""
|
||||
img_item = ImageContent(data=b"abc", mime_type="image/jpeg")
|
||||
assistant_message = ChatMessageContent(role=AuthorRole.ASSISTANT, items=[img_item])
|
||||
|
||||
with pytest.raises(ServiceInvalidRequestError) as exc:
|
||||
MESSAGE_CONVERTERS[AuthorRole.ASSISTANT](assistant_message)
|
||||
|
||||
assert "Image content is not supported in an assistant message." in str(exc.value)
|
||||
|
||||
|
||||
def test_format_assistant_message_unsupported_type() -> None:
|
||||
"""Test assistant message with unsupported item content type."""
|
||||
func_res_item = FunctionResultContent(id="res1", function_name="some_function", result="some_result")
|
||||
assistant_message = ChatMessageContent(role=AuthorRole.ASSISTANT, items=[func_res_item])
|
||||
|
||||
with pytest.raises(ServiceInvalidRequestError) as exc:
|
||||
MESSAGE_CONVERTERS[AuthorRole.ASSISTANT](assistant_message)
|
||||
assert "Unsupported content type in an assistant message:" in str(exc.value)
|
||||
|
||||
|
||||
def test_format_tool_message_text() -> None:
|
||||
"""Test tool message with text content."""
|
||||
text_item = TextContent(text="Some text")
|
||||
tool_message = ChatMessageContent(role=AuthorRole.TOOL, items=[text_item])
|
||||
|
||||
formatted = MESSAGE_CONVERTERS[AuthorRole.TOOL](tool_message)
|
||||
assert formatted["role"] == "user" # note that for a tool message, role set to 'user'
|
||||
assert formatted["content"] == [{"text": "Some text"}]
|
||||
|
||||
|
||||
def test_format_tool_message_function_result() -> None:
|
||||
"""Test tool message with function result content."""
|
||||
func_result_item = FunctionResultContent(id="res_id", function_name="test_function", result="some result")
|
||||
tool_message = ChatMessageContent(role=AuthorRole.TOOL, items=[func_result_item])
|
||||
|
||||
formatted = MESSAGE_CONVERTERS[AuthorRole.TOOL](tool_message)
|
||||
assert formatted["role"] == "user"
|
||||
content = formatted["content"][0]
|
||||
assert content.get("toolResult")
|
||||
assert content["toolResult"]["toolUseId"] == "res_id"
|
||||
assert content["toolResult"]["content"] == [{"text": "some result"}]
|
||||
|
||||
|
||||
def test_format_tool_message_image_raises() -> None:
|
||||
"""Test tool message with image content raises an error."""
|
||||
img_item = ImageContent(data=b"xyz", mime_type="image/jpeg")
|
||||
tool_message = ChatMessageContent(role=AuthorRole.TOOL, items=[img_item])
|
||||
|
||||
with pytest.raises(ServiceInvalidRequestError) as exc:
|
||||
MESSAGE_CONVERTERS[AuthorRole.TOOL](tool_message)
|
||||
assert "Image content is not supported in a tool message." in str(exc.value)
|
||||
|
||||
|
||||
def test_finish_reason_from_bedrock_to_semantic_kernel_stop() -> None:
|
||||
"""Test that 'stop_sequence' maps to FinishReason.STOP"""
|
||||
reason = finish_reason_from_bedrock_to_semantic_kernel("stop_sequence")
|
||||
assert reason == FinishReason.STOP
|
||||
|
||||
reason = finish_reason_from_bedrock_to_semantic_kernel("end_turn")
|
||||
assert reason == FinishReason.STOP
|
||||
|
||||
|
||||
def test_finish_reason_from_bedrock_to_semantic_kernel_length() -> None:
|
||||
"""Test that 'max_tokens' maps to FinishReason.LENGTH"""
|
||||
reason = finish_reason_from_bedrock_to_semantic_kernel("max_tokens")
|
||||
assert reason == FinishReason.LENGTH
|
||||
|
||||
|
||||
def test_finish_reason_from_bedrock_to_semantic_kernel_content_filtered() -> None:
|
||||
"""Test that 'content_filtered' maps to FinishReason.CONTENT_FILTER"""
|
||||
reason = finish_reason_from_bedrock_to_semantic_kernel("content_filtered")
|
||||
assert reason == FinishReason.CONTENT_FILTER
|
||||
|
||||
|
||||
def test_finish_reason_from_bedrock_to_semantic_kernel_tool_use() -> None:
|
||||
"""Test that 'tool_use' maps to FinishReason.TOOL_CALLS"""
|
||||
reason = finish_reason_from_bedrock_to_semantic_kernel("tool_use")
|
||||
assert reason == FinishReason.TOOL_CALLS
|
||||
|
||||
|
||||
def test_finish_reason_from_bedrock_to_semantic_kernel_unknown() -> None:
|
||||
"""Test that unknown finish reason returns None"""
|
||||
reason = finish_reason_from_bedrock_to_semantic_kernel("something_unknown")
|
||||
assert reason is None
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_bedrock_settings() -> BedrockChatPromptExecutionSettings:
|
||||
"""Helper fixture for BedrockChatPromptExecutionSettings."""
|
||||
return BedrockChatPromptExecutionSettings()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_function_choice_config() -> FunctionCallChoiceConfiguration:
|
||||
"""Helper fixture for a sample FunctionCallChoiceConfiguration."""
|
||||
|
||||
# We'll create mock kernel functions with metadata
|
||||
mock_func_1 = MagicMock()
|
||||
mock_func_1.fully_qualified_name = "plugin-function1"
|
||||
mock_func_1.description = "Function 1 description"
|
||||
|
||||
param1 = MagicMock()
|
||||
param1.name = "param1"
|
||||
param1.schema_data = {"type": "string"}
|
||||
param1.is_required = True
|
||||
|
||||
param2 = MagicMock()
|
||||
param2.name = "param2"
|
||||
param2.schema_data = {"type": "integer"}
|
||||
param2.is_required = False
|
||||
|
||||
mock_func_1.parameters = [
|
||||
param1,
|
||||
param2,
|
||||
]
|
||||
mock_func_2 = MagicMock()
|
||||
mock_func_2.fully_qualified_name = "plugin-function2"
|
||||
mock_func_2.description = "Function 2 description"
|
||||
mock_func_2.parameters = []
|
||||
|
||||
config = FunctionCallChoiceConfiguration()
|
||||
config.available_functions = [mock_func_1, mock_func_2]
|
||||
|
||||
return config
|
||||
|
||||
|
||||
def test_update_settings_from_function_choice_configuration_none_type(
|
||||
mock_function_choice_config, mock_bedrock_settings
|
||||
) -> None:
|
||||
"""Test that if the FunctionChoiceType is NONE it doesn't modify settings."""
|
||||
update_settings_from_function_choice_configuration(
|
||||
mock_function_choice_config, mock_bedrock_settings, FunctionChoiceType.NONE
|
||||
)
|
||||
assert mock_bedrock_settings.tool_choice is None
|
||||
assert mock_bedrock_settings.tools is None
|
||||
|
||||
|
||||
def test_update_settings_from_function_choice_configuration_auto_two_tools(
|
||||
mock_function_choice_config, mock_bedrock_settings
|
||||
) -> None:
|
||||
"""Test that AUTO sets tool_choice to {"auto": {}} and sets tools list"""
|
||||
update_settings_from_function_choice_configuration(
|
||||
mock_function_choice_config, mock_bedrock_settings, FunctionChoiceType.AUTO
|
||||
)
|
||||
assert mock_bedrock_settings.tool_choice == {"auto": {}}
|
||||
assert len(mock_bedrock_settings.tools) == 2
|
||||
# Validate structure of first tool
|
||||
tool_spec_1 = mock_bedrock_settings.tools[0].get("toolSpec")
|
||||
assert tool_spec_1["name"] == "plugin-function1"
|
||||
assert tool_spec_1["description"] == "Function 1 description"
|
||||
|
||||
|
||||
def test_update_settings_from_function_choice_configuration_required_many(
|
||||
mock_function_choice_config, mock_bedrock_settings
|
||||
) -> None:
|
||||
"""Test that REQUIRED with more than one function sets tool_choice to {"any": {}}."""
|
||||
update_settings_from_function_choice_configuration(
|
||||
mock_function_choice_config, mock_bedrock_settings, FunctionChoiceType.REQUIRED
|
||||
)
|
||||
assert mock_bedrock_settings.tool_choice == {"any": {}}
|
||||
assert len(mock_bedrock_settings.tools) == 2
|
||||
|
||||
|
||||
def test_update_settings_from_function_choice_configuration_required_one(mock_bedrock_settings) -> None:
|
||||
"""Test that REQUIRED with a single function picks "tool" with that function name."""
|
||||
single_func = MagicMock()
|
||||
single_func.fully_qualified_name = "plugin-function"
|
||||
single_func.description = "Only function"
|
||||
single_func.parameters = []
|
||||
|
||||
config = FunctionCallChoiceConfiguration()
|
||||
config.available_functions = [single_func]
|
||||
|
||||
update_settings_from_function_choice_configuration(config, mock_bedrock_settings, FunctionChoiceType.REQUIRED)
|
||||
assert mock_bedrock_settings.tool_choice == {"tool": {"name": "plugin-function"}}
|
||||
assert len(mock_bedrock_settings.tools) == 1
|
||||
assert mock_bedrock_settings.tools[0]["toolSpec"]["name"] == "plugin-function"
|
||||
@@ -0,0 +1,324 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import json
|
||||
from functools import reduce
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import boto3
|
||||
import pytest
|
||||
|
||||
from semantic_kernel.connectors.ai.bedrock.bedrock_prompt_execution_settings import BedrockTextPromptExecutionSettings
|
||||
from semantic_kernel.connectors.ai.bedrock.services.bedrock_text_completion import BedrockTextCompletion
|
||||
from semantic_kernel.connectors.ai.bedrock.services.model_provider.bedrock_model_provider import (
|
||||
BedrockModelProvider,
|
||||
get_text_completion_request_body,
|
||||
)
|
||||
from semantic_kernel.contents.streaming_text_content import StreamingTextContent
|
||||
from semantic_kernel.contents.text_content import TextContent
|
||||
from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError
|
||||
from tests.unit.connectors.ai.bedrock.conftest import MockBedrockClient, MockBedrockRuntimeClient
|
||||
|
||||
# region init
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_bedrock_text_completion_init(mock_client, bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Completion service"""
|
||||
bedrock_text_completion = BedrockTextCompletion()
|
||||
|
||||
assert bedrock_text_completion.ai_model_id == bedrock_unit_test_env["BEDROCK_TEXT_MODEL_ID"]
|
||||
assert bedrock_text_completion.service_id == bedrock_unit_test_env["BEDROCK_TEXT_MODEL_ID"]
|
||||
|
||||
assert bedrock_text_completion.bedrock_model_provider == BedrockModelProvider(
|
||||
bedrock_unit_test_env["BEDROCK_MODEL_PROVIDER"]
|
||||
)
|
||||
assert bedrock_text_completion.bedrock_client is not None
|
||||
assert bedrock_text_completion.bedrock_runtime_client is not None
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_bedrock_text_completion_init_model_id_override(mock_client, bedrock_unit_test_env, model_id) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Completion service"""
|
||||
bedrock_text_completion = BedrockTextCompletion(model_id=model_id)
|
||||
|
||||
assert bedrock_text_completion.ai_model_id == model_id
|
||||
assert bedrock_text_completion.service_id == model_id
|
||||
|
||||
assert bedrock_text_completion.bedrock_client is not None
|
||||
assert bedrock_text_completion.bedrock_runtime_client is not None
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_bedrock_text_completion_init_custom_service_id(mock_client, bedrock_unit_test_env, service_id) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Completion service"""
|
||||
bedrock_text_completion = BedrockTextCompletion(service_id=service_id)
|
||||
|
||||
assert bedrock_text_completion.service_id == service_id
|
||||
|
||||
assert bedrock_text_completion.bedrock_client is not None
|
||||
assert bedrock_text_completion.bedrock_runtime_client is not None
|
||||
|
||||
|
||||
def test_bedrock_text_completion_init_custom_clients(bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Completion service"""
|
||||
bedrock_text_completion = BedrockTextCompletion(
|
||||
runtime_client=MockBedrockRuntimeClient(),
|
||||
client=MockBedrockClient(),
|
||||
)
|
||||
|
||||
assert isinstance(bedrock_text_completion.bedrock_client, MockBedrockClient)
|
||||
assert isinstance(bedrock_text_completion.bedrock_runtime_client, MockBedrockRuntimeClient)
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_bedrock_text_completion_init_custom_client(mock_client, bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Completion service"""
|
||||
bedrock_text_completion = BedrockTextCompletion(
|
||||
client=MockBedrockClient(),
|
||||
)
|
||||
|
||||
assert isinstance(bedrock_text_completion.bedrock_client, MockBedrockClient)
|
||||
assert bedrock_text_completion.bedrock_runtime_client is not None
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_bedrock_text_completion_init_custom_runtime_client(mock_client, bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Completion service"""
|
||||
bedrock_text_completion = BedrockTextCompletion(
|
||||
runtime_client=MockBedrockRuntimeClient(),
|
||||
)
|
||||
|
||||
assert bedrock_text_completion.bedrock_client is not None
|
||||
assert isinstance(bedrock_text_completion.bedrock_runtime_client, MockBedrockRuntimeClient)
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_bedrock_text_completion_init_custom_bedrock_model_provider(mock_client, bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Completion service"""
|
||||
bedrock_text_completion = BedrockTextCompletion(
|
||||
model_provider=BedrockModelProvider.AMAZON,
|
||||
)
|
||||
|
||||
assert bedrock_text_completion.bedrock_model_provider == BedrockModelProvider.AMAZON
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["BEDROCK_TEXT_MODEL_ID"]], indirect=True)
|
||||
def test_bedrock_text_completion_client_init_with_empty_model_id(bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Completion service with empty model id"""
|
||||
with pytest.raises(ServiceInitializationError, match="The Amazon Bedrock Text Model ID is missing."):
|
||||
BedrockTextCompletion(env_file_path="fake_env_file_path.env")
|
||||
|
||||
|
||||
def test_bedrock_text_completion_client_init_invalid_settings(bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Completion service with invalid settings"""
|
||||
with pytest.raises(
|
||||
ServiceInitializationError, match="Failed to initialize the Amazon Bedrock Text Completion Service."
|
||||
):
|
||||
BedrockTextCompletion(model_id=123) # Model ID must be a string
|
||||
|
||||
|
||||
def test_bedrock_text_completion_client_init_invalid_model_provider(bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Completion service with invalid settings"""
|
||||
with pytest.raises(
|
||||
ServiceInitializationError, match="Failed to initialize the Amazon Bedrock Text Completion Service."
|
||||
):
|
||||
BedrockTextCompletion(model_provider="invalid_provider")
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_prompt_execution_settings_class(mock_client, bedrock_unit_test_env) -> None:
|
||||
"""Test getting prompt execution settings class"""
|
||||
bedrock_completion_client = BedrockTextCompletion()
|
||||
assert bedrock_completion_client.get_prompt_execution_settings_class() == BedrockTextPromptExecutionSettings
|
||||
|
||||
|
||||
# endregion
|
||||
|
||||
|
||||
# region text completion
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
# These are fake model ids with the supported prefixes
|
||||
"model_id",
|
||||
[
|
||||
"amazon.titan",
|
||||
"anthropic.claude",
|
||||
"cohere.command",
|
||||
"ai21.jamba",
|
||||
"meta.llama",
|
||||
"mistral.ai",
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
async def test_bedrock_text_completion(
|
||||
model_id,
|
||||
mock_bedrock_text_completion_response,
|
||||
output_text,
|
||||
) -> None:
|
||||
"""Test Amazon Bedrock Text Completion complete method"""
|
||||
with patch.object(
|
||||
MockBedrockRuntimeClient, "invoke_model", return_value=mock_bedrock_text_completion_response
|
||||
) as mock_model_invoke:
|
||||
# Setup
|
||||
bedrock_text_completion = BedrockTextCompletion(
|
||||
model_id=model_id,
|
||||
runtime_client=MockBedrockRuntimeClient(),
|
||||
client=MockBedrockClient(),
|
||||
)
|
||||
|
||||
# Act
|
||||
settings = BedrockTextPromptExecutionSettings()
|
||||
response = await bedrock_text_completion.get_text_contents("Hello!", settings=settings)
|
||||
|
||||
# Assert
|
||||
mock_model_invoke.assert_called_once_with(
|
||||
body=json.dumps(get_text_completion_request_body(model_id, "Hello!", settings)),
|
||||
modelId=model_id,
|
||||
accept="application/json",
|
||||
contentType="application/json",
|
||||
)
|
||||
|
||||
assert isinstance(response, list)
|
||||
assert len(response) == 1
|
||||
assert isinstance(response[0], TextContent)
|
||||
assert response[0].ai_model_id == model_id
|
||||
assert response[0].text == output_text
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
# These are fake model ids with the supported prefixes
|
||||
"model_id",
|
||||
[
|
||||
"arn:aws:bedrock:us-east-1:972143716085:application-inference-profile/123456",
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
async def test_bedrock_text_completion_with_application_inference_profile(
|
||||
model_id,
|
||||
mock_bedrock_text_completion_response,
|
||||
output_text,
|
||||
model_provider,
|
||||
) -> None:
|
||||
"""Test Amazon Bedrock Text Completion complete method"""
|
||||
with patch.object(
|
||||
MockBedrockRuntimeClient,
|
||||
"invoke_model",
|
||||
return_value=mock_bedrock_text_completion_response,
|
||||
) as mock_model_invoke:
|
||||
# Setup
|
||||
bedrock_text_completion = BedrockTextCompletion(
|
||||
model_id=model_id,
|
||||
runtime_client=MockBedrockRuntimeClient(),
|
||||
client=MockBedrockClient(),
|
||||
model_provider=model_provider,
|
||||
)
|
||||
|
||||
# Act
|
||||
settings = BedrockTextPromptExecutionSettings()
|
||||
await bedrock_text_completion.get_text_contents("Hello!", settings=settings)
|
||||
|
||||
# Assert
|
||||
mock_model_invoke.assert_called_once_with(
|
||||
body=json.dumps(get_text_completion_request_body(model_id, "Hello!", settings, model_provider)),
|
||||
modelId=model_id,
|
||||
accept="application/json",
|
||||
contentType="application/json",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
# These are fake model ids with the supported prefixes
|
||||
"model_id",
|
||||
[
|
||||
"amazon.titan",
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
async def test_bedrock_streaming_text_completion(
|
||||
model_id,
|
||||
mock_bedrock_streaming_text_completion_response,
|
||||
output_text,
|
||||
) -> None:
|
||||
"""Test Amazon Bedrock Text Completion complete method"""
|
||||
with patch.object(
|
||||
MockBedrockRuntimeClient,
|
||||
"invoke_model_with_response_stream",
|
||||
return_value=mock_bedrock_streaming_text_completion_response,
|
||||
) as mock_invoke_model_with_response_stream:
|
||||
# Setup
|
||||
bedrock_text_completion = BedrockTextCompletion(
|
||||
model_id=model_id,
|
||||
runtime_client=MockBedrockRuntimeClient(),
|
||||
client=MockBedrockClient(),
|
||||
)
|
||||
|
||||
# Act
|
||||
settings = BedrockTextPromptExecutionSettings()
|
||||
chunks: list[StreamingTextContent] = []
|
||||
async for streaming_responses in bedrock_text_completion.get_streaming_text_contents(
|
||||
"Hello!", settings=settings
|
||||
):
|
||||
chunks.extend(streaming_responses)
|
||||
response = reduce(lambda p, r: p + r, chunks)
|
||||
|
||||
# Assert
|
||||
mock_invoke_model_with_response_stream.assert_called_once_with(
|
||||
body=json.dumps(get_text_completion_request_body(model_id, "Hello!", settings)),
|
||||
modelId=model_id,
|
||||
accept="application/json",
|
||||
contentType="application/json",
|
||||
)
|
||||
assert isinstance(response, StreamingTextContent)
|
||||
assert response.ai_model_id == model_id
|
||||
assert response.text == output_text
|
||||
assert response.choice_index == 0
|
||||
assert isinstance(response.inner_content, list)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
# These are fake model ids with the supported prefixes
|
||||
"model_id",
|
||||
[
|
||||
"arn:aws:bedrock:us-east-1:972143716085:application-inference-profile/123456",
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
async def test_bedrock_streaming_text_completion_with_application_inference_profile(
|
||||
model_id,
|
||||
mock_bedrock_streaming_text_completion_response,
|
||||
output_text,
|
||||
model_provider,
|
||||
) -> None:
|
||||
"""Test Amazon Bedrock Chat Completion complete method"""
|
||||
with patch.object(
|
||||
MockBedrockRuntimeClient,
|
||||
"invoke_model_with_response_stream",
|
||||
return_value=mock_bedrock_streaming_text_completion_response,
|
||||
) as mock_invoke_model_with_response_stream:
|
||||
# Setup
|
||||
bedrock_text_completion = BedrockTextCompletion(
|
||||
model_id=model_id,
|
||||
runtime_client=MockBedrockRuntimeClient(),
|
||||
client=MockBedrockClient(),
|
||||
model_provider=model_provider,
|
||||
)
|
||||
|
||||
# Act
|
||||
settings = BedrockTextPromptExecutionSettings()
|
||||
chunks: list[StreamingTextContent] = []
|
||||
async for streaming_responses in bedrock_text_completion.get_streaming_text_contents(
|
||||
"Hello!", settings=settings
|
||||
):
|
||||
chunks.extend(streaming_responses)
|
||||
|
||||
# Assert
|
||||
mock_invoke_model_with_response_stream.assert_called_once_with(
|
||||
body=json.dumps(get_text_completion_request_body(model_id, "Hello!", settings, model_provider)),
|
||||
modelId=model_id,
|
||||
accept="application/json",
|
||||
contentType="application/json",
|
||||
)
|
||||
|
||||
|
||||
# endregion
|
||||
+235
@@ -0,0 +1,235 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from unittest.mock import ANY, Mock, patch
|
||||
|
||||
import boto3
|
||||
import pytest
|
||||
|
||||
from semantic_kernel.connectors.ai.bedrock.bedrock_prompt_execution_settings import (
|
||||
BedrockEmbeddingPromptExecutionSettings,
|
||||
)
|
||||
from semantic_kernel.connectors.ai.bedrock.services.bedrock_text_embedding import BedrockTextEmbedding
|
||||
from semantic_kernel.connectors.ai.bedrock.services.model_provider.bedrock_model_provider import BedrockModelProvider
|
||||
from semantic_kernel.exceptions.service_exceptions import ServiceInitializationError, ServiceInvalidResponseError
|
||||
from tests.unit.connectors.ai.bedrock.conftest import MockBedrockClient, MockBedrockRuntimeClient
|
||||
|
||||
# region init
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_bedrock_text_embedding_init(mock_client, bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Embedding service"""
|
||||
bedrock_text_embedding = BedrockTextEmbedding()
|
||||
|
||||
assert bedrock_text_embedding.ai_model_id == bedrock_unit_test_env["BEDROCK_EMBEDDING_MODEL_ID"]
|
||||
assert bedrock_text_embedding.service_id == bedrock_unit_test_env["BEDROCK_EMBEDDING_MODEL_ID"]
|
||||
|
||||
assert bedrock_text_embedding.bedrock_model_provider == BedrockModelProvider(
|
||||
bedrock_unit_test_env["BEDROCK_MODEL_PROVIDER"]
|
||||
)
|
||||
assert bedrock_text_embedding.bedrock_client is not None
|
||||
assert bedrock_text_embedding.bedrock_runtime_client is not None
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_bedrock_text_embedding_init_model_id_override(mock_client, bedrock_unit_test_env, model_id) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Embedding service"""
|
||||
bedrock_text_embedding = BedrockTextEmbedding(model_id=model_id)
|
||||
|
||||
assert bedrock_text_embedding.ai_model_id == model_id
|
||||
assert bedrock_text_embedding.service_id == model_id
|
||||
|
||||
assert bedrock_text_embedding.bedrock_client is not None
|
||||
assert bedrock_text_embedding.bedrock_runtime_client is not None
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_bedrock_text_embedding_init_custom_service_id(mock_client, bedrock_unit_test_env, service_id) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Embedding service"""
|
||||
bedrock_text_embedding = BedrockTextEmbedding(service_id=service_id)
|
||||
|
||||
assert bedrock_text_embedding.service_id == service_id
|
||||
|
||||
assert bedrock_text_embedding.bedrock_client is not None
|
||||
assert bedrock_text_embedding.bedrock_runtime_client is not None
|
||||
|
||||
|
||||
def test_bedrock_text_embedding_init_custom_clients(bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Embedding service"""
|
||||
bedrock_text_embedding = BedrockTextEmbedding(
|
||||
runtime_client=MockBedrockRuntimeClient(),
|
||||
client=MockBedrockClient(),
|
||||
)
|
||||
|
||||
assert isinstance(bedrock_text_embedding.bedrock_client, MockBedrockClient)
|
||||
assert isinstance(bedrock_text_embedding.bedrock_runtime_client, MockBedrockRuntimeClient)
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_bedrock_text_embedding_init_custom_client(mock_client, bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Embedding service"""
|
||||
bedrock_text_embedding = BedrockTextEmbedding(
|
||||
client=MockBedrockClient(),
|
||||
)
|
||||
|
||||
assert isinstance(bedrock_text_embedding.bedrock_client, MockBedrockClient)
|
||||
assert bedrock_text_embedding.bedrock_runtime_client is not None
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_bedrock_text_embedding_init_custom_runtime_client(mock_client, bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Embedding service"""
|
||||
bedrock_text_embedding = BedrockTextEmbedding(
|
||||
runtime_client=MockBedrockRuntimeClient(),
|
||||
)
|
||||
|
||||
assert bedrock_text_embedding.bedrock_client is not None
|
||||
assert isinstance(bedrock_text_embedding.bedrock_runtime_client, MockBedrockRuntimeClient)
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_bedrock_text_embedding_init_custom_bedrock_model_provider(mock_client, bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Embedding service"""
|
||||
bedrock_text_embedding = BedrockTextEmbedding(
|
||||
model_provider=BedrockModelProvider.AMAZON,
|
||||
)
|
||||
|
||||
assert bedrock_text_embedding.bedrock_model_provider == BedrockModelProvider.AMAZON
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exclude_list", [["BEDROCK_EMBEDDING_MODEL_ID"]], indirect=True)
|
||||
def test_bedrock_text_embedding_client_init_with_empty_model_id(bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Embedding service with empty model id"""
|
||||
with pytest.raises(ServiceInitializationError, match="The Amazon Bedrock Text Embedding Model ID is missing."):
|
||||
BedrockTextEmbedding(env_file_path="fake_env_file_path.env")
|
||||
|
||||
|
||||
def test_bedrock_text_embedding_client_init_invalid_settings(bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Embedding service with invalid settings"""
|
||||
with pytest.raises(
|
||||
ServiceInitializationError, match="Failed to initialize the Amazon Bedrock Text Embedding Service."
|
||||
):
|
||||
BedrockTextEmbedding(model_id=123) # Model ID must be a string
|
||||
|
||||
|
||||
def test_bedrock_text_embedding_client_init_invalid_model_provider(bedrock_unit_test_env) -> None:
|
||||
"""Test initialization of Amazon Bedrock Text Embedding service with invalid settings"""
|
||||
with pytest.raises(
|
||||
ServiceInitializationError, match="Failed to initialize the Amazon Bedrock Text Embedding Service."
|
||||
):
|
||||
BedrockTextEmbedding(model_provider="invalid_provider")
|
||||
|
||||
|
||||
@patch.object(boto3, "client", return_value=Mock())
|
||||
def test_prompt_execution_settings_class(mock_client, bedrock_unit_test_env) -> None:
|
||||
"""Test getting prompt execution settings class"""
|
||||
bedrock_completion_client = BedrockTextEmbedding()
|
||||
assert bedrock_completion_client.get_prompt_execution_settings_class() == BedrockEmbeddingPromptExecutionSettings
|
||||
|
||||
|
||||
# endregion
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
# These are fake model ids with the supported prefixes
|
||||
"model_id",
|
||||
[
|
||||
"amazon.titan",
|
||||
"cohere.command",
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
async def test_bedrock_text_embedding(model_id, mock_bedrock_text_embedding_response) -> None:
|
||||
"""Test Bedrock text embedding generation"""
|
||||
with patch.object(
|
||||
MockBedrockRuntimeClient, "invoke_model", return_value=mock_bedrock_text_embedding_response
|
||||
) as mock_model_invoke:
|
||||
# Setup
|
||||
bedrock_text_embedding = BedrockTextEmbedding(
|
||||
model_id=model_id,
|
||||
runtime_client=MockBedrockRuntimeClient(),
|
||||
client=MockBedrockClient(),
|
||||
)
|
||||
|
||||
# Act
|
||||
settings = BedrockEmbeddingPromptExecutionSettings()
|
||||
response = await bedrock_text_embedding.generate_embeddings(["hello", "world"], settings)
|
||||
|
||||
# Assert
|
||||
mock_model_invoke.assert_called_with(
|
||||
body=ANY,
|
||||
modelId=model_id,
|
||||
accept="application/json",
|
||||
contentType="application/json",
|
||||
)
|
||||
assert mock_model_invoke.call_count == 2
|
||||
|
||||
assert len(response) == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
# These are fake model ids with the supported prefixes
|
||||
"model_id",
|
||||
[
|
||||
"arn:aws:bedrock:us-east-1:972143716085:application-inference-profile/123456",
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
async def test_bedrock_text_embedding_with_application_inference_profile(
|
||||
model_id,
|
||||
mock_bedrock_text_embedding_response,
|
||||
model_provider,
|
||||
) -> None:
|
||||
"""Test Bedrock text embedding generation"""
|
||||
with patch.object(
|
||||
MockBedrockRuntimeClient, "invoke_model", return_value=mock_bedrock_text_embedding_response
|
||||
) as mock_model_invoke:
|
||||
# Setup
|
||||
bedrock_text_embedding = BedrockTextEmbedding(
|
||||
model_id=model_id,
|
||||
runtime_client=MockBedrockRuntimeClient(),
|
||||
client=MockBedrockClient(),
|
||||
model_provider=BedrockModelProvider.AMAZON,
|
||||
)
|
||||
|
||||
# Act
|
||||
settings = BedrockEmbeddingPromptExecutionSettings()
|
||||
response = await bedrock_text_embedding.generate_embeddings(["hello", "world"], settings)
|
||||
|
||||
# Assert
|
||||
mock_model_invoke.assert_called_with(
|
||||
body=ANY,
|
||||
modelId=model_id,
|
||||
accept="application/json",
|
||||
contentType="application/json",
|
||||
)
|
||||
assert mock_model_invoke.call_count == 2
|
||||
|
||||
assert len(response) == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
# These are fake model ids with the supported prefixes
|
||||
"model_id",
|
||||
[
|
||||
"amazon.titan",
|
||||
"cohere.command",
|
||||
],
|
||||
indirect=True,
|
||||
)
|
||||
async def test_bedrock_text_embedding_with_invalid_response(
|
||||
model_id, mock_bedrock_text_embedding_invalid_response
|
||||
) -> None:
|
||||
"""Test Bedrock text embedding generation with invalid response"""
|
||||
with patch.object(
|
||||
MockBedrockRuntimeClient, "invoke_model", return_value=mock_bedrock_text_embedding_invalid_response
|
||||
):
|
||||
# Setup
|
||||
bedrock_text_embedding = BedrockTextEmbedding(
|
||||
model_id=model_id,
|
||||
runtime_client=MockBedrockRuntimeClient(),
|
||||
client=MockBedrockClient(),
|
||||
)
|
||||
|
||||
with pytest.raises(ServiceInvalidResponseError):
|
||||
await bedrock_text_embedding.generate_embeddings(["hello", "world"])
|
||||
Reference in New Issue
Block a user